use randn on GPU instead of cpu

nianticlabs · Jan 29, 2022 · e33aef8 · e33aef8
1 parent 93b9ebb
commit e33aef8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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: