import torch 
import argparse
from assembler import get_config, assembler
from data import  CelebADataModule
from pytorch_lightning import Trainer
import torchvision.transforms as transforms

# Load configs
parser = argparse.ArgumentParser(description='Generic runner for VAE models')
parser.add_argument('--config',  '-c',
                    dest="filename",
                    metavar='FILE',
                    help =  'path to the config file',
                    default='configs/vae.yaml')

args = parser.parse_args()
config = get_config(args.filename)
gan = assembler(config)

# Load data
trans = []
trans.append(transforms.RandomHorizontalFlip())
if config["exp_params"]["crop_size"] > 0:
  trans.append(transforms.CenterCrop(config["exp_params"]["crop_size"]))
trans.append(transforms.Resize(config["exp_params"]["img_size"]))
trans.append(transforms.ToTensor())
transform = transforms.Compose(trans)

dm = CelebADataModule(data_dir=config["exp_params"]["data_path"],
                                target_type='attr',
                                train_transform=transform,
                                val_transform=transform,
                                download=True,
                                batch_size=config["exp_params"]["batch_size"])

# Run Training Loop
trainer= Trainer(gpus = config["trainer_params"]["gpus"],
        max_epochs = config["trainer_params"]["max_epochs"])
trainer.fit(gan, datamodule=dm)
torch.save(gan.state_dict(), f"{gan.name}_celeba.ckpt")
@@ -73,18 +73,29 @@
 r%0Afrom d
-at
+ownload_celeb
 a import
@@ -95,33 +95,33 @@
  import 
- C
+c
 eleb
-ADataModule
+a_dataloader
 %0Afrom py
@@ -608,145 +608,30 @@
 ata%0A
-trans = %5B%5D%0Atrans.append(transforms.RandomHorizontalFlip())%0Aif config%5B%22exp_params%22%5D%5B%22crop_size%22%5D %3E 0:%0A  trans.append(transforms.CenterCrop
+dm = celeba_dataloader
 (con
@@ -653,442 +653,135 @@
 %22%5D%5B%22
-crop
+batch
 _size%22%5D
-))%0Atrans.append(transforms.Resize(config%5B%22exp_params%22%5D%5B%22img_size%22%5D))%0Atrans.append(transforms.ToTensor())%0Atransform = transforms.Compose(trans)%0A%0Adm = CelebADataModule(data_dir=config%5B%22exp_params%22%5D%5B%22data_path%22%5D,%0A                                target_type='attr',%0A                                train_transform=transform,%0A                                val_transform=transform,%0A                                download=True
+,%0A                       config%5B%22exp_params%22%5D%5B%22img_size%22%5D,%0A                       config%5B%22exp_params%22%5D%5B%22crop_size%22%5D
 ,%0A  
@@ -806,27 +806,8 @@
     
-        batch_size=
 conf
@@ -828,23 +828,23 @@
 %22%5D%5B%22
-batch_size
+data_path
 %22%5D)%0A%0A
+%0A
 # Ru
@@ -1065,8 +1065,9 @@
 a.ckpt%22)
+%0A

update run.py to use celeba dataloader