Merge pull request #413 from d4l3k/master

use randn on GPU instead of cpu
nianticlabs · Jan 30, 2022 · b676244 · b676244
2 parents 93b9ebb + e33aef8
commit b676244
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: