Move to gpu and add checkpoint images to attributes

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

train_vqvae.py

@@ -34,7 +34,7 @@ def train(args):
     np.random.seed(seed)
 
     model = VQVAE(
-        im_channels=dataset_config["im_channels"], model_config=autoencoder_config)
+        im_channels=dataset_config["im_channels"], model_config=autoencoder_config).to(device)
 
     data_loader = create_dataloader(dataset_config["im_path"])
 
@@ -48,7 +48,7 @@ def train(args):
 
     # LPIPS model
     lpips_model = LPIPS().eval().to(device)
-    discriminator = Discriminator(im_channels=dataset_config["im_channels"])
+    discriminator = Discriminator(im_channels=dataset_config["im_channels"]).to(device)
 
     optimizer_d = Adam(discriminator.parameters(),
                        lr=train_config["autoencoder_lr"], betas=(0.5, 0.999))
@@ -76,9 +76,8 @@ def train(args):
             steps = 0
         
             # Model output with losses
-            print("Sending batch to model")
+            im_tensor = im_tensor.to(device)
             model_output = model(im_tensor)
-            print(model_output)
             output, z, quatize_losses = model_output