Resolve run time errors

dataset/celeba.py

@@ -40,9 +40,10 @@ def create_dataloader(parquet_dir, batch_size=32, shuffle=True, num_workers=4):
 
 
 # Test the dataloader
-dataloader = create_dataloader('dataset')
-for i, batch in enumerate(dataloader):
-    if i == 0:
-        print(batch[0].shape, batch[1])
-    else:
-        break
+# dataloader = create_dataloader('dataset')
+# for i, batch in enumerate(dataloader):
+#     for img_tensor in batch:
+#         if i == 0 :
+#             print(img_tensor)
+#         else:
+#             break

models/blocks.py

@@ -109,7 +109,7 @@ class DownBlock(nn.Module):
             # Resnet Block
             resnet_input = out
             out = self.resnet_conv_first[i](out)
-            if self.t_emb is not None:
+            if self.t_emb_dim is not None:
                 out = out + self.time_embd_layers[i](t_emb)[:, :, None, None]
             out = self.resnet_conv_second[i](out)
             out = out + self.residual_input_conv[i](resnet_input)

train_vqvae.py

@@ -12,6 +12,7 @@ from dataset.celeba import create_dataloader
 from torchvision.utils import make_grid
 import yaml
 import numpy as np
+from tqdm import tqdm
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
@@ -23,7 +24,6 @@ def train(args):
         except yaml.YAMLError as e:
             print(e)
 
-    print(config)
 
     autoencoder_config = config["autoencoder_params"]
     train_config = config["train_config"]
@@ -73,9 +73,10 @@ def train(args):
         optimizer_d.zero_grad()
         optimizer_g.zero_grad()
 
-        for im_tensor, _ in data_loader:
+        for image in tqdm(data_loader):
             steps = 0
-            im_tensor = im_tensor.float().to(device)
+            print(type(image[0]))
+            im_tensor = image[0]
 
             # Model output with losses
             model_output = model(im_tensor)