Model is running without issues. next step is to work on input density.

Author: Hananeh Oliaei

src/electrai/chk.pth

@@ -59,14 +59,12 @@ def rand_rotate(self, data_lst):
             return [rotate(rotate(rotate(d))) for d in data_lst]
 
     def __getitem__(self, idx):
-        rho1 = torch.tensor(
-            np.load(self.data[idx]), dtype=torch.float32)
-        size = np.loadtxt(self.data_gs[idx], dtype=int)
+        rho1 = torch.tensor(self.data[idx], dtype=torch.float32)
+        size = torch.tensor(self.data_gs[idx], dtype=int)
         rho1 = rho1.reshape(1, *size)
 
-        rho2 = torch.tensor(
-            np.load(self.label[idx]), dtype=torch.float32)
-        size = np.loadtxt(self.label_gs[idx], dtype=int)
+        rho2 = torch.tensor(self.label[idx], dtype=torch.float32)
+        size = torch.tensor(self.label_gs[idx], dtype=int)
         rho2 = rho2.reshape(1, *size)
 
         if self.da:

src/electrai/chk_1.pth

@@ -105,11 +105,8 @@ def loss_fn_sum(output, target):
             return loss
 
     optimizer.zero_grad()
-    # print(dataloader)
-    # for batch, (X, y) in enumerate(dataloader):
-    for batch, cont in enumerate(dataloader):
+    for batch, (X, y) in enumerate(dataloader):
         print('batch: ', batch)
-        print('cont: ', cont)
         # print('X: ', X.shape)
         # print('y: ', y.shape)
         X, y = X.to(cfg.device), y.to(cfg.device)