minor fixes

README.md

@@ -31,9 +31,9 @@ You can create an anaconda environment called `conv_onet` using
 conda env create -f environment.yaml
 conda activate conv_onet
 ```
-**Note**: you might need to install **torch-scatter** mannually following [the official instruction](https://github.com/rusty1s/pytorch_scatter#pytorch-140):
+**Note**: you might need to (re-)install **torch-scatter** mannually following [the official instruction](https://github.com/rusty1s/pytorch_scatter#pytorch-140):
 ```
-pip install torch-scatter==2.0.4 -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html
+pip install torch-scatter==2.0.4 --force-reinstall -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html
 ```
 
 Next, compile the extension modules.

generate.py

@@ -40,7 +40,7 @@
 # Model
 model = config.get_model(cfg, device=device, dataset=dataset)
 
-checkpoint_io = CheckpointIO(out_dir, model=model)
+checkpoint_io = CheckpointIO(out_dir, model=model, is_cuda=is_cuda)
 checkpoint_io.load(cfg['test']['model_file'])
 
 # Generator

src/checkpoints.py

@@ -12,8 +12,9 @@ class CheckpointIO(object):
     Args:
         checkpoint_dir (str): path where checkpoints are saved
     '''
-    def __init__(self, checkpoint_dir='./chkpts', **kwargs):
+    def __init__(self, checkpoint_dir='./chkpts', is_cuda=True, **kwargs):
         self.module_dict = kwargs
+        self.is_cuda = is_cuda
         self.checkpoint_dir = checkpoint_dir
         if not os.path.exists(checkpoint_dir):
             os.makedirs(checkpoint_dir)
@@ -75,7 +76,10 @@ def load_url(self, url):
         '''
         print(url)
         print('=> Loading checkpoint from url...')
-        state_dict = model_zoo.load_url(url, progress=True)
+        if self.is_cuda:
+            state_dict = model_zoo.load_url(url, progress=True)
+        else:
+            state_dict = model_zoo.load_url(url, progress=True, map_location=torch.device('cpu'))
         scalars = self.parse_state_dict(state_dict)
         return scalars