Allocate tensors directly on target device

src/boltz/model/models/boltz1.py

@@ -508,7 +508,7 @@ def training_step(self, batch: dict[str, Tensor], batch_idx: int) -> Tensor:
             )
         else:
             confidence_loss_dict = {
-                "loss": torch.tensor(0.0).to(batch["token_index"].device),
+                "loss": torch.tensor(0.0, device=batch["token_index"].device),
                 "loss_breakdown": {},
             }
 

src/boltz/model/modules/confidence.py

@@ -427,8 +427,10 @@ def forward(
         pred_distogram_prob = nn.functional.softmax(
             pred_distogram_logits, dim=-1
         ).repeat_interleave(multiplicity, 0)
-        contacts = torch.zeros((1, 1, 1, 64), dtype=pred_distogram_prob.dtype).to(
-            pred_distogram_prob.device
+        contacts = torch.zeros(
+            (1, 1, 1, 64),
+            dtype=pred_distogram_prob.dtype,
+            device=pred_distogram_prob.device,
         )
         contacts[:, :, :, :20] = 1.0
         prob_contact = (pred_distogram_prob * contacts).sum(-1)

src/boltz/model/modules/confidencev2.py

@@ -434,8 +434,10 @@ def forward(
         pred_distogram_prob = nn.functional.softmax(
             pred_distogram_logits, dim=-1
         ).repeat_interleave(multiplicity, 0)
-        contacts = torch.zeros((1, 1, 1, 64), dtype=pred_distogram_prob.dtype).to(
-            pred_distogram_prob.device
+        contacts = torch.zeros(
+            (1, 1, 1, 64),
+            dtype=pred_distogram_prob.dtype,
+            device=pred_distogram_prob.device,
         )
         contacts[:, :, :, :20] = 1.0
         prob_contact = (pred_distogram_prob * contacts).sum(-1)

src/boltz/model/modules/diffusion.py

@@ -531,9 +531,13 @@ def sample(
 
             with torch.no_grad():
                 atom_coords_denoised = torch.zeros_like(atom_coords_noisy)
-                token_a = torch.zeros(token_repr_shape).to(atom_coords_noisy)
+                token_a = torch.zeros(
+                    token_repr_shape,
+                    device=atom_coords_noisy.device,
+                    dtype=atom_coords_noisy.dtype,
+                )
 
-                sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device)
+                sample_ids = torch.arange(multiplicity, device=atom_coords_noisy.device)
                 sample_ids_chunks = sample_ids.chunk(
                     multiplicity % max_parallel_samples + 1
                 )

src/boltz/model/modules/diffusionv2.py

@@ -379,7 +379,7 @@ def sample(
 
             with torch.no_grad():
                 atom_coords_denoised = torch.zeros_like(atom_coords_noisy)
-                sample_ids = torch.arange(multiplicity).to(atom_coords_noisy.device)
+                sample_ids = torch.arange(multiplicity, device=atom_coords_noisy.device)
                 sample_ids_chunks = sample_ids.chunk(
                     multiplicity % max_parallel_samples + 1
                 )

src/boltz/model/modules/encoders.py

@@ -510,7 +510,7 @@ def forward(
             # only here the multiplicity kicks in because we use the different positions r
             q = q.repeat_interleave(multiplicity, 0)
             r_input = torch.cat(
-                [r, torch.zeros((B * multiplicity, N, 7)).to(r)],
+                [r, torch.zeros((B * multiplicity, N, 7), device=r.device, dtype=r.dtype)],
                 dim=-1,
             )
             r_to_q = self.r_to_q_trans(r_input)