Repository: Stonesjtu/pytorch_memlab

Issue:
Variable 'tensor_names' referenced before assignment
https://github.com/Stonesjtu/pytorch_memlab/blob/ec9a72fc302981ddc3ee56d6e16694610d646c36/pytorch_memlab/mem_reporter.py#L37

The variable tensor_names is referenced before assignment if no model is passed into the MemReporter. Need to move it into the above if statement.

Comments:

Patch:
diff --git a/.travis.yml b/.travis.yml
index 4c040a2..c36f90a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -14,18 +14,22 @@ matrix:
     env:
       IPYTHON_VERSION='7'
       PANDAS_VERSION='1'
+  - python: '3.8'
+    env:
+      IPYTHON_VERSION='7'
+      PANDAS_VERSION='1'
 install:
 - pip install IPython==$IPYTHON_VERSION
 - pip install pandas==$PANDAS_VERSION
-- python setup.py install
-- pip install -r requirements.txt
+- pip install .[test]
 script:
 - python -c 'import pytorch_memlab'
+- pytest test/test_mem_reporter.py
 deploy:
   provider: pypi
   user: yoursky
   on:
     tags: true
-    python: 3.7
+    python: 3.8
   password:
     secure: P0hSqy9eGKVG/0Cu8yE1+I+V0gIUbc+B7srJTsyJPC4+COv7BqIrt1NG+3+vWdocJLuYiTbf3zpUBozdvHH0XDV8Ki0rjsQOm607Fe+d30JCaL3sXLJOe8qwbizx9nmTS04AORc7nyzn4Tc3H10zDfwxL5uk73H6mwaoMlh2v2sizB8G6Ce4kH9JXWarKahKPPSdt2u0o533Bm5/rfNzxLifGmG7o2OYIXmUbeC54f4nhBCihtc+sjKR54qfopH/Gnpl1fnFK2Q5aMOM2yUj7mrjNLgDeAECigKGZpn/uuYyy/dSJKvPFoHFP52HS6x+9/aVUUgwEhqyYZ3tFJgkyeLPVnuP5Pv7wQyZXbIBchPwljswgjxI/+8ANRM6WMUhBbnOUQPPd/6AmW6xdaJf2l0461jGqhxAUGRKpn/odFIEly3TcizIzHjkOqPbS4xkKgN7s40ai8ZFLIUXmbqa7r/dScdu7qjvRYIn+obTCIq0lR3gTZLNfHHBCYFOcLD0anlDakONaiY4++xzDw88ancLQhN5L5rsQge4QNdZS8s88gbPtei+3DfnGsUnYWWplAHdxJ+A8/CrtBrtM18E3mOuwSdKbCwd54YmL9E+KcRcL/WRpedWjNybiBCDvRoO0iw3+2EnkrTLpsuTtiXAu+oe2h1XoMyYbeVjaVVsN44=
diff --git a/pytorch_memlab/line_profiler/line_profiler.py b/pytorch_memlab/line_profiler/line_profiler.py
index 04c1fa2..7615ef1 100644
--- a/pytorch_memlab/line_profiler/line_profiler.py
+++ b/pytorch_memlab/line_profiler/line_profiler.py
@@ -6,7 +6,7 @@ import torch
 from .line_records import LineRecords
 
 # Seaborn's `muted` color cycle
-DEFAULT_COLUMNS = ['active_bytes.all.peak', 'reserved_bytes.all.peak']
+DEFAULT_COLUMNS = ('active_bytes.all.peak', 'reserved_bytes.all.peak')
 
 
 class LineProfiler:
@@ -88,8 +88,9 @@ class LineProfiler:
         try:
             torch.cuda.empty_cache()
             self._reset_cuda_stats()
-        except AssertionError as e:
-            print('Could not reset CUDA stats and cache: ' + str(e))
+        # Pytorch-1.7.0 raises AttributeError while <1.6.0 raises AssertionError
+        except (AssertionError, AttributeError) as error:
+            print('Could not reset CUDA stats and cache: ' + str(error))
 
         self.register_callback()
 
diff --git a/pytorch_memlab/mem_reporter.py b/pytorch_memlab/mem_reporter.py
index 6aff270..c2f3800 100644
--- a/pytorch_memlab/mem_reporter.py
+++ b/pytorch_memlab/mem_reporter.py
@@ -27,13 +27,15 @@ class MemReporter():
         self.device_tensor_stat = {}
         # to numbering the unknown tensors
         self.name_idx = 0
+
+        tensor_names = defaultdict(list)
         if model is not None:
             assert isinstance(model, torch.nn.Module)
             # for model with tying weight, multiple parameters may share
             # the same underlying tensor
-            tensor_names = defaultdict(list)
             for name, param in model.named_parameters():
                 tensor_names[param].append(name)
+
         for param, name in tensor_names.items():
             self.tensor_name[id(param)] = '+'.join(name)
 


Test patch:
diff --git a/test/test_mem_reporter.py b/test/test_mem_reporter.py
index 971e8f1..927d7fe 100644
--- a/test/test_mem_reporter.py
+++ b/test/test_mem_reporter.py
@@ -7,8 +7,8 @@ import pytest
 concentrate_mode = False
 
 def test_reporter():
-    linear = torch.nn.Linear(1024, 1024).cuda()
-    inp = torch.Tensor(512, 1024).cuda()
+    linear = torch.nn.Linear(1024, 1024)
+    inp = torch.Tensor(512, 1024)
     reporter = MemReporter(linear)
 
     out = linear(inp*(inp+3)).mean()
@@ -17,16 +17,27 @@ def test_reporter():
 
     reporter.report()
 
+def test_reporter_without_model():
+    linear = torch.nn.Linear(1024, 1024)
+    inp = torch.Tensor(512, 1024)
+    reporter = MemReporter()
+
+    out = linear(inp*(inp+3)).mean()
+    reporter.report()
+    out.backward()
+
+    reporter.report()
+
 @pytest.mark.skipif(concentrate_mode, reason='concentrate')
 def test_reporter_tie_weight():
-    linear = torch.nn.Linear(1024, 1024).cuda()
-    linear_2 = torch.nn.Linear(1024, 1024).cuda()
+    linear = torch.nn.Linear(1024, 1024)
+    linear_2 = torch.nn.Linear(1024, 1024)
     linear_2.weight = linear.weight
     container = torch.nn.Sequential(
         linear, linear_2
     )
     reporter = MemReporter(container)
-    inp = torch.Tensor(512, 1024).cuda()
+    inp = torch.Tensor(512, 1024)
 
     out = container(inp).mean()
     out.backward()
@@ -34,6 +45,7 @@ def test_reporter_tie_weight():
     reporter = MemReporter(container)
     reporter.report()
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
 @pytest.mark.skipif(concentrate_mode, reason='concentrate')
 def test_reporter_LSTM():
     lstm = torch.nn.LSTM(256, 256, num_layers=1).cuda()
@@ -45,6 +57,7 @@ def test_reporter_LSTM():
     reporter = MemReporter(lstm)
     reporter.report()
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
 @pytest.mark.skipif(concentrate_mode, reason='concentrate')
 def test_reporter_device():
     lstm_cpu = torch.nn.LSTM(256, 256)