small bugs

Author: Ubuntu

src/aging_gan/data.py

@@ -132,7 +132,7 @@ def prepare_dataset(
     # randomness
     train_transform = T.Compose(
         [
-            T.ToPILImage(),
+            # T.ToPILImage(),
             T.RandomHorizontalFlip(),
             T.Resize((img_size + 50, img_size + 50), antialias=True),
             T.RandomCrop(img_size),

src/aging_gan/train.py

@@ -94,7 +94,7 @@ def parse_args() -> argparse.Namespace:
     p.add_argument(
         "--steps_for_logging_metrics",
         type=int,
-        default=1,
+        default=50,
         help="Print training metrics after certain batch steps.",
     )
     p.add_argument(

src/aging_gan/utils.py

@@ -103,9 +103,18 @@ def generate_and_save_samples(
     device: torch.device,
     num_samples: int = 8,
 ):
-    # grab one batch
-    inputs, _ = next(iter(val_loader))
-    inputs = inputs.to(device)[:num_samples]
+    # grab batches until num_samples   
+    collected = []
+    for imgs, _ in val_loader:
+        collected.append(imgs)
+        if sum(b.size(0) for b in collected) >= num_samples:
+            break
+
+    if not collected:
+        raise ValueError("Validation loader is empty.")
+    
+    inputs = torch.cat(collected, dim=0)[:num_samples].to(device)
+    
     with torch.no_grad():
         outputs = generator(inputs)