Updates

brlp_lite.py

@@ -32,10 +32,12 @@
 # }
 
 import os
+os.environ["PYTORCH_WEIGHTS_ONLY"] = "False"
 from typing import Optional, Union
 import pandas as pd
 import argparse
 import numpy as np
+
 import warnings
 import torch
 import torch.nn as nn
@@ -43,7 +45,7 @@ from torch import Tensor
 from torch.optim.optimizer import Optimizer
 from torch.nn import L1Loss
 from torch.utils.data import DataLoader
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 from torch.amp import GradScaler
 
 from generative.networks.nets import (
@@ -57,8 +59,12 @@ from monai import transforms
 from monai.utils import set_determinism
 from monai.data.meta_tensor import MetaTensor
 import torch.serialization
-
+from numpy.core.multiarray import _reconstruct
+from numpy import ndarray, dtype
+torch.serialization.add_safe_globals([_reconstruct])
 torch.serialization.add_safe_globals([MetaTensor])
+torch.serialization.add_safe_globals([ndarray])
+torch.serialization.add_safe_globals([dtype])
 
 from tqdm import tqdm
 import matplotlib.pyplot as plt
@@ -381,6 +387,22 @@ def train(
     train_df = dataset_df[dataset_df.split == 'train']
     trainset = get_dataset_from_pd(train_df, transforms_fn, cache_dir)
 
+    print(f"[DEBUG] Using cache_dir={cache_dir}")
+    print(f"[DEBUG] trainset length={len(trainset)}")
+
+    try:
+        sample_debug = trainset[0]  # Force a transform on the first record
+        print("[DEBUG] Successfully loaded sample 0 from trainset.")
+    except Exception as e:
+        print("[DEBUG] Error loading sample 0:", e)
+
+    import glob
+
+    hashfiles = glob.glob(os.path.join(cache_dir, "*.pt"))
+    print(f"[DEBUG] Found {len(hashfiles)} cached .pt files in {cache_dir}")
+    if hashfiles:
+        print("[DEBUG] Example cache file:", hashfiles[0])
+
     train_loader = DataLoader(
         dataset=trainset,
         num_workers=num_workers,
@@ -440,13 +462,14 @@ def train(
     total_counter = 0
 
     for epoch in range(n_epochs):
+        print(f"[DEBUG] Starting epoch {epoch}/{n_epochs-1}")
         autoencoder.train()
         progress_bar = tqdm(enumerate(train_loader), total=len(train_loader))
-        progress_bar.set_description(f'Epoch {epoch + 1}/{n_epochs}')
+        progress_bar.set_description(f'Epoch {epoch}')
 
         for step, batch in progress_bar:
             # Generator Training
-            with autocast(enabled=True):
+            with autocast(device, enabled=True):
                 images = batch["image"].to(device)
                 reconstruction, z_mu, z_sigma = autoencoder(images)
 
@@ -462,7 +485,7 @@ def train(
             gradacc_g.step(loss_g, step)
 
             # Discriminator Training
-            with autocast(enabled=True):
+            with autocast(device, enabled=True):
                 logits_fake = discriminator(reconstruction.contiguous().detach())[-1]
                 d_loss_fake = adv_loss_fn(logits_fake, target_is_real=False, for_discriminator=True)
                 logits_real = discriminator(images.contiguous().detach())[-1]
@@ -604,4 +627,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()