From e33aef868c68c5f565c42a20170f8b4ea1daa566 Mon Sep 17 00:00:00 2001
From: Tristan Rice <rice@fn.lc>
Date: Fri, 28 Jan 2022 17:24:44 -0800
Subject: [PATCH] use randn on GPU instead of cpu

---
 trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/trainer.py b/trainer.py
index a726dad64..f370603dc 100644
--- a/trainer.py
+++ b/trainer.py
@@ -466,7 +466,7 @@ def compute_losses(self, inputs, outputs):
             if not self.opt.disable_automasking:
                 # add random numbers to break ties
                 identity_reprojection_loss += torch.randn(
-                    identity_reprojection_loss.shape).cuda() * 0.00001
+                    identity_reprojection_loss.shape, device=self.device) * 0.00001
 
                 combined = torch.cat((identity_reprojection_loss, reprojection_loss), dim=1)
             else: