Update prediction.py

prediction.py

@@ -55,11 +55,11 @@ def run_image_prediction(
 
     # Sample from model using provided sequence and nucleus image
     _, _, _, predicted_threshold, predicted_heatmap = model.celle.sample(
-        text=sequence,
-        condition=nucleus_image,
+        text=sequence.to(device),
+        condition=nucleus_image.to(device),
         timesteps=1,
         temperature=1,
-        progress=True,
+        progress=False,
     )
 
     # Move predicted_threshold and predicted_heatmap to CPU and select first element of batch