Removing GPU dependency

Learning-and-Intelligent-Systems · Dec 10, 2020 · cfaf61a · cfaf61a
1 parent 7dd40a6
commit cfaf61a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,2 @@
+*.pyc
+
diff --git a/model/template.py b/model/template.py
@@ -38,7 +38,8 @@ def get_random_points(self, shape, device="gpu0"):
         Return Tensor of Size [x, 3, x ... x]
         """
         assert shape[1] == 3, "shape should have 3 in dim 1"
-        rand_grid = torch.cuda.FloatTensor(shape).to(device).float()
+        FloatTensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
+        rand_grid = FloatTensor(shape).to(device).float()
         rand_grid.data.normal_(0, 1)
         rand_grid = rand_grid / torch.sqrt(torch.sum(rand_grid ** 2, dim=1, keepdim=True))
         return Variable(rand_grid)