Update gunpowder_trainer.py to swap to(device) and float()

janelia-cellmap · Oct 24, 2024 · ee7c762 · ee7c762
1 parent 4152f7c
commit ee7c762
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dacapo/experiments/trainers/gunpowder_trainer.py b/dacapo/experiments/trainers/gunpowder_trainer.py
@@ -309,7 +309,7 @@ def iterate(self, num_iterations, model, optimizer, device):
                 param.grad = None
 
             t_start_prediction = time.time()
-            predicted = model.forward(torch.as_tensor(raw[raw.roi]).to(device).float())
+            predicted = model.forward(torch.as_tensor(raw[raw.roi]).float().to(device))
             predicted.retain_grad()
             loss = self._loss.compute(
                 predicted,