Upload layers.py

pgm/layers.py

@@ -173,8 +173,8 @@ class ArgMaxGumbelMax(Transform):
         
         def sample_gumbel(shape, eps=1e-20):
             U = torch.rand(shape)
-            
-            U = U.to(torch.device('cuda:1'))
+            U = U.to(torch.device('cpu'))
+            # U = U.to(torch.device('cuda:1'))
             return -torch.log(-torch.log(U + eps) + eps)
         def gumbel_softmax_sample(logits, temperature):
             y = logits + sample_gumbel(logits.shape)