Upload folder using huggingface_hub

gram_smol_latent/train.py

@@ -183,9 +183,12 @@ class VQGANModel(flax.struct.PyTreeNode):
         
                 return loss
 
+            #latents = results_dict["latents"]
 
-            B, H, W, C = reconstructed_images.shape
-            reshaped_features = reconstructed_images.reshape(B, -1, C)
+            #B, H, W, C = reconstructed_images.shape
+            #reshaped_features = reconstructed_images.reshape(B, -1, C)
+            B, H, W, C = result_dict["latents"].shape
+            reshaped_features = result_dict["latents"].reshape(B, -1, C)
             batched_loss_fn = jax.vmap(calculate_covariance_loss_single, in_axes=0)
             per_image_losses = batched_loss_fn(reconstructed_images)