Catch RuntimeError in torch.cuda.reset_peak_memory_stats (#26)

The AttributeError was fixed in pytorch/pytorch#48406
Stonesjtu · Dec 6, 2020 · 004895f · 004895f
1 parent 654cbb8
commit 004895f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/pytorch_memlab/line_profiler/line_profiler.py b/pytorch_memlab/line_profiler/line_profiler.py
@@ -88,8 +88,9 @@ def enable(self):
         try:
             torch.cuda.empty_cache()
             self._reset_cuda_stats()
-        # Pytorch-1.7.0 raises AttributeError while <1.6.0 raises AssertionError
-        except (AssertionError, AttributeError) as error:
+        # What error is raised depends on PyTorch version:
+        # latest raises RuntimeError, 1.7.0 raises AttributeError, <1.7.0 raises AssertionError
+        except (AssertionError, AttributeError, RuntimeError) as error:
             print('Could not reset CUDA stats and cache: ' + str(error))
 
         self.register_callback()