Fix test_no_grad

kakaobrain · Jun 10, 2019 · b1c1a35 · b1c1a35
1 parent 38458bf
commit b1c1a35
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_gpipe.py b/tests/test_gpipe.py
@@ -226,7 +226,7 @@ def hook(module, input, outputs):
         nonlocal latent
         latent = output
 
-    partition = list(model.partitions())[0]
+    partition = model.partitions[0]
     partition.register_forward_hook(hook)
 
     with torch.no_grad():