From ee7c762725ea6d86663b30e612bf285621522d32 Mon Sep 17 00:00:00 2001
From: Peter Sobolewski <76622105+psobolewskiPhD@users.noreply.github.com>
Date: Thu, 24 Oct 2024 11:54:12 +0200
Subject: [PATCH] Update gunpowder_trainer.py to swap to(device) and float()

---
 dacapo/experiments/trainers/gunpowder_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dacapo/experiments/trainers/gunpowder_trainer.py b/dacapo/experiments/trainers/gunpowder_trainer.py
index 4b5649e77..104c5fa9c 100644
--- 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,