VAE Debugged and Running

lib/modules/utils.py

@@ -6,8 +6,6 @@ from torch import nn
 from torch import functional as F
 from torch.utils.data import DataLoader
 
-from lib.objects.map import MapStorage
-
 import pytorch_lightning as pl
 
 
@@ -27,10 +25,11 @@ class Flatten(nn.Module):
             print(e)
             return -1
 
-    def __init__(self, in_shape, to=(-1, )):
+    def __init__(self, in_shape, to=-1):
+        assert isinstance(to, int) or isinstance(to, tuple)
         super(Flatten, self).__init__()
         self.in_shape = in_shape
-        self.to = to
+        self.to = (to,) if isinstance(to, int) else to
 
     def forward(self, x):
         return x.view(x.size(0), *self.to)
@@ -107,7 +106,7 @@ class LightningBaseModule(pl.LightningModule, ABC):
         # Data loading
         # =============================================================================
         # Map Object
-        self.map_storage = MapStorage(self.hparams.data_param.map_root)
+        # self.map_storage = MapStorage(self.hparams.data_param.map_root)
 
     def size(self):
         return self.shape