Use 'cpu' for testing balance_by_time

kakaobrain · Nov 9, 2019 · ff3bab8 · ff3bab8
1 parent 981d918
commit ff3bab8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_balance.py b/tests/test_balance.py
@@ -167,7 +167,7 @@ def forward(self, x):
     assert not model.training
 
     sample = torch.rand(1)
-    balance_by_time(1, model, sample)
+    balance_by_time(1, model, sample, device='cpu')
 
     assert not model.training
 
@@ -189,7 +189,7 @@ def forward(self, a_b):
 
     model = nn.Sequential(Twin(), Add())
     sample = torch.rand(1, requires_grad=True)
-    balance_by_time(1, model, sample)
+    balance_by_time(1, model, sample, device='cpu')
 
 
 @skip_if_no_cuda