Fix end-of-epoch with record_stream

definitelynotmcarilli · definitelynotmcarilli · commit 90e5b05a2d2f · 2019-06-20T14:12:26.000-07:00
diff --git a/examples/imagenet/main_amp.py b/examples/imagenet/main_amp.py
@@ -300,8 +300,10 @@ def next(self):
         torch.cuda.current_stream().wait_stream(self.stream)
         input = self.next_input
         target = self.next_target
-        input.record_stream(torch.cuda.current_stream())
-        target.record_stream(torch.cuda.current_stream())
+        if input is not None:
+            input.record_stream(torch.cuda.current_stream())
+        if target is not None:
+            target.record_stream(torch.cuda.current_stream())
         self.preload()
         return input, target
 
