fix pytest

Makefile

@@ -20,11 +20,8 @@ clean:
 	@python3 -Bc "import shutil, pathlib; \
 		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('__pycache__')]; \
 		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.ipynb_checkpoints')]; \
-		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.monai-cache')]; \
-		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.mypy_cache')]; \
-		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.ruff_cache')]; \
-		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.pytest_cache')]"
+		[shutil.rmtree(p) for p in pathlib.Path('.').rglob('.monai-cache')];"
 
 # Updated 'check' to clean before running (optional)
 # This ensures you are testing from a "blank slate"
-check: format lint typecheck clean
+check: format lint typecheck test clean

run_cspca.py

@@ -21,7 +21,6 @@ def main_worker(args):
     cache_dir_path = Path(os.path.join(args.logdir, "cache"))
 
     if args.mode == "train":
-
         checkpoint = torch.load(args.checkpoint_pirads, weights_only=False, map_location="cpu")
         mil_model.load_state_dict(checkpoint["state_dict"])
         mil_model = mil_model.to(args.device)
@@ -64,7 +63,6 @@ def main_worker(args):
         if cache_dir_path.exists() and cache_dir_path.is_dir():
             shutil.rmtree(cache_dir_path)
 
-
     cspca_model = CSPCAModel(backbone=mil_model).to(args.device)
     checkpt = torch.load(args.checkpoint_cspca, map_location="cpu")
     cspca_model.load_state_dict(checkpt["state_dict"])
@@ -92,7 +90,6 @@ def main_worker(args):
     get_metrics(metrics_dict)
 
 
-
 def parse_args():
     parser = argparse.ArgumentParser(
         description="Multiple Instance Learning (MIL) for csPCa risk prediction."

src/data/data_loader.py

@@ -26,6 +26,7 @@ from .custom_transforms import (
     NormalizeIntensity_customd,
 )
 
+
 class DummyMILDataset(torch.utils.data.Dataset):
     def __init__(self, args, num_samples=8):
         self.num_samples = num_samples
@@ -43,13 +44,16 @@ class DummyMILDataset(torch.utils.data.Dataset):
             item = {
                 # Shape: (Channels=3, Depth, H, W) based on your Transposed(indices=(0, 3, 1, 2))
                 "image": torch.randn(3, self.args.depth, self.args.tile_size, self.args.tile_size),
-                "label": torch.tensor(label_value, dtype=torch.float32)
+                "label": torch.tensor(label_value, dtype=torch.float32),
             }
             if self.args.use_heatmap:
-                item["final_heatmap"] = torch.randn(1, self.args.depth, self.args.tile_size, self.args.tile_size)
+                item["final_heatmap"] = torch.randn(
+                    1, self.args.depth, self.args.tile_size, self.args.tile_size
+                )
             bag.append(item)
         return bag
 
+
 def list_data_collate(batch: list):
     """
     Combine instances from a list of dicts into a single dict, by stacking them along first dim
@@ -130,18 +134,16 @@ def data_transform(args: argparse.Namespace) -> Transform:
 def get_dataloader(
     args: argparse.Namespace, split: Literal["train", "test"]
 ) -> torch.utils.data.DataLoader:
-
     if args.dry_run:
         print(f"🛠️  DRY RUN: Creating synthetic {split} dataloader...")
         dummy_ds = DummyMILDataset(args, num_samples=args.batch_size * 2)
         return torch.utils.data.DataLoader(
             dummy_ds,
             batch_size=args.batch_size,
-            collate_fn=list_data_collate, # Uses your custom stacking logic
-            num_workers=0 # Keep it simple for dry run
+            collate_fn=list_data_collate,  # Uses your custom stacking logic
+            num_workers=0,  # Keep it simple for dry run
         )
 
-
     data_list = load_decathlon_datalist(
         data_list_file_path=args.dataset_json,
         data_list_key=split,

temp.ipynb

@@ -77,6 +77,249 @@
     "from monai.data import Dataset, load_decathlon_datalist, ITKReader, NumpyReader, PersistentDataset"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "bc433898",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import subprocess\n",
+    "import sys\n",
+    "from pathlib import Path\n",
+    "import torch\n",
+    "import pytest\n",
+    "import argparse\n",
+    "from src.train.train_pirads import get_attention_scores\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "f1c90aff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batch_size = 2\n",
+    "num_patches = 4\n",
+    "\n",
+    "# Sample 0: Target = 3 (Cancer), Sample 1: Target = 0 (PI-RADS 2)\n",
+    "data = torch.randn(batch_size, num_patches, 1, 8, 8)\n",
+    "target = torch.tensor([3.0, 0.0])\n",
+    "\n",
+    "# Create heatmaps: Sample 0 has one \"hot\" patch\n",
+    "heatmap = torch.zeros(batch_size, num_patches, 1, 8, 8)\n",
+    "heatmap[0, 0] = 10.0 # High attention on patch 0 for the first sample\n",
+    "heatmap[0, 3] = 2.0 \n",
+    "heatmap[1, 2] = 5.0  # Should be overridden by PI-RADS 2 logic anyway\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "80cb444f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def mock_args():\n",
+    "    # Mocking argparse for the device\n",
+    "    args = argparse.Namespace()\n",
+    "    args.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "    return args"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "id": "6528fd4d",
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "AssertionError",
+     "evalue": "",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[41], line 23\u001b[0m\n\u001b[1;32m     21\u001b[0m idx \u001b[38;5;241m=\u001b[39m (shuffled_images[\u001b[38;5;241m0\u001b[39m, :, \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m] \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m5.0\u001b[39m)\u001b[38;5;241m.\u001b[39mnonzero(as_tuple\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)[\u001b[38;5;241m0\u001b[39m]\n\u001b[1;32m     22\u001b[0m \u001b[38;5;66;03m# The attention score at that same index should be the maximum\u001b[39;00m\n\u001b[0;32m---> 23\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m att_labels[\u001b[38;5;241m0\u001b[39m, idx] \u001b[38;5;241m==\u001b[39m att_labels[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mmedian()\n",
+      "\u001b[0;31mAssertionError\u001b[0m: "
+     ]
+    }
+   ],
+   "source": [
+    "num_patches = 10\n",
+    "\n",
+    "# Distinct data per patch: [0, 1, 2, 3...]\n",
+    "data = torch.arange(num_patches).view(1, num_patches, 1, 1, 1).float()\n",
+    "target = torch.tensor([3.0])\n",
+    "\n",
+    "# Heatmap matches the data indices so we can track the \"label\"\n",
+    "heatmap = torch.arange(num_patches).view(1, num_patches, 1, 1, 1).float()\n",
+    "\n",
+    "att_labels, shuffled_images = get_attention_scores(data, target, heatmap, mock_args)\n",
+    "\n",
+    "\n",
+    "idx= (shuffled_images[0, :, 0, 0, 0] == 9.0).nonzero(as_tuple=True)[0]\n",
+    "# The attention score at that same index should be the maximum\n",
+    "assert att_labels[0, idx] == att_labels[0].max()\n",
+    "\n",
+    "idx = (shuffled_images[0, :, 0, 0, 0] == 0.0).nonzero(as_tuple=True)[0]\n",
+    "# The attention score at that same index should be the maximum\n",
+    "assert att_labels[0, idx] == att_labels[0].min()\n",
+    "\n",
+    "idx = (shuffled_images[0, :, 0, 0, 0] == 5.0).nonzero(as_tuple=True)[0]\n",
+    "# The attention score at that same index should be the maximum\n",
+    "assert att_labels[0, idx] == att_labels[0].median()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "90f5acab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import subprocess\n",
+    "import sys\n",
+    "from pathlib import Path\n",
+    "import torch\n",
+    "import pytest\n",
+    "import argparse\n",
+    "from src.train.train_pirads import get_attention_scores\n",
+    "import monai\n",
+    "from monai.transforms import Transform\n",
+    "from src.data.custom_transforms import NormalizeIntensity_custom"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "e3a2dc6c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "img = torch.zeros((2, 4, 4), dtype=torch.float32)\n",
+    "mask = torch.zeros((1, 4, 4), dtype=torch.float32)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "98a500df",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([[[0., 0., 0., 0.],\n",
+       "         [0., 0., 0., 0.],\n",
+       "         [0., 0., 0., 0.],\n",
+       "         [0., 0., 0., 0.]],\n",
+       "\n",
+       "        [[0., 0., 0., 0.],\n",
+       "         [0., 0., 0., 0.],\n",
+       "         [0., 0., 0., 0.],\n",
+       "         [0., 0., 0., 0.]]])"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "img"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "c9974f43",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "img[0, :, :] = 100.0  # Background\n",
+    "img[0, 0, 0] = 10.0   # Masked pixel 1\n",
+    "img[0, 0, 1] = 20.0   # Masked pixel 2\n",
+    "\n",
+    "# --- Channel 1 Setup ---\n",
+    "# Inside mask: Values [2, 4]\n",
+    "# Outside mask: Value 50\n",
+    "img[1, :, :] = 50.0   # Background\n",
+    "img[1, 0, 0] = 2.0    # Masked pixel 1\n",
+    "img[1, 0, 1] = 4.0    # Masked pixel 2\n",
+    "\n",
+    "# --- Mask Setup ---\n",
+    "# Selects only the top-left two pixels (0,0) and (0,1)\n",
+    "mask[0, 0, 0] = 1\n",
+    "mask[0, 0, 1] = 1\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "eb910fda",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data = torch.rand(1, 10, 10)\n",
+    "mask = torch.randint(0, 2, (1, 10, 10)).float()\n",
+    "normalizer = NormalizeIntensity_custom(nonzero=False, channel_wise=True)\n",
+    "out = normalizer(data, mask)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "923341a3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "masked = data[mask != 0]\n",
+    "mean_ = torch.mean(masked.float())\n",
+    "std_ = torch.std(masked.float(), unbiased=False)\n",
+    "\n",
+    "epsilon = 1e-8\n",
+    "normalized_data = (data - mean_) / (std_ + epsilon)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "e844cde1",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([ 1.4106, -0.1975,  0.3907,  1.2870, -0.7974, -1.2061,  0.7028,  1.2778,\n",
+       "         0.4667, -0.3361, -0.7842, -1.6296, -1.2037,  1.3582, -0.5648, -0.3055,\n",
+       "        -0.3313,  0.0328, -1.0675,  0.6328, -0.2215, -1.3372,  0.5165,  1.9302,\n",
+       "         0.8875,  0.6793,  0.5553,  0.4335,  0.6390, -1.3707,  1.6053,  1.8626,\n",
+       "        -0.3923,  0.2319,  0.3911, -0.4683, -1.1255, -1.6464, -0.2123, -0.5415,\n",
+       "         0.1401, -0.2822,  1.5019, -0.5117, -1.6047, -0.2322, -1.3080,  0.0130,\n",
+       "         1.8028,  0.5602, -1.6317])"
+      ]
+     },
+     "execution_count": 23,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "masked"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "a9a20f58",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.testing.assert_close(out, normalized_data)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 2,

tests/test_run.py

@@ -1,139 +1,169 @@
-import subprocess
-import sys
-from pathlib import Path
+import argparse
 
+import pytest
+import torch
 
-def test_run_pirads_training():
-    """
-    Test that run_cspca.py runs without crashing using an existing YAML config.
-    """
+from src.data.custom_transforms import NormalizeIntensity_custom
+from src.data.data_loader import get_dataloader
+from src.model.cspca_model import CSPCAModel
+from src.model.mil import MILModel3D
+from src.train import train_cspca, train_pirads
+from src.train.train_pirads import get_attention_scores
 
-    # Path to your run_pirads.py script
-    repo_root = Path(__file__).parent.parent
-    script_path = repo_root / "run_pirads.py"
 
-    # Path to your existing config.yaml
-    config_path = repo_root / "config" / "config_pirads_train.yaml"  # adjust this path
+@pytest.fixture
+def mock_args():
+    # Mocking argparse for the device
+    args = argparse.Namespace()
+    args.device = "cuda" if torch.cuda.is_available() else "cpu"
+    return args
 
-    # Make sure the file exists
-    assert config_path.exists(), f"Config file not found: {config_path}"
 
-    # Run the script with the config
-    result = subprocess.run(
-        [
-            sys.executable,
-            str(script_path),
-            "--mode",
-            "train",
-            "--config",
-            str(config_path),
-            "--dry_run",
-        ],
-        capture_output=True,
-        text=True,
-    )
+def test_get_attention_scores_logic(mock_args):
+    # Setup: 2 samples, 4 patches, images of size 8x8
+    batch_size = 2
+    num_patches = 4
 
-    # Check that it ran without errors
-    assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
+    # Sample 0: Target = 3 (Cancer), Sample 1: Target = 0 (PI-RADS 2)
+    data = torch.randn(batch_size, num_patches, 1, 8, 8)
+    target = torch.tensor([3.0, 0.0])
 
+    # Create heatmaps: Sample 0 has one "hot" patch
+    heatmap = torch.zeros(batch_size, num_patches, 1, 8, 8)
+    heatmap[0, 0] = 10.0  # High attention on patch 0 for the first sample
+    heatmap[1, :] = 5.0  # Should be overridden by PI-RADS 2 logic anyway
+
+    att_labels, shuffled_images = get_attention_scores(data, target, heatmap, mock_args)
+
+    # --- TEST 1: Normalization ---
+    sums = att_labels.sum(dim=1)
+    torch.testing.assert_close(sums, torch.ones(batch_size).to(mock_args.device))
+
+    # --- TEST 2: PI-RADS 2 Uniformity ---
+    pirads_2_scores = att_labels[1]
+    expected_uniform = torch.ones(num_patches).to(mock_args.device) / num_patches
+    torch.testing.assert_close(pirads_2_scores, expected_uniform)
+
+    # --- TEST 4: Output Shapes ---
+    assert att_labels.shape == (batch_size, num_patches)
+    assert shuffled_images.shape == data.shape
 
-def test_run_pirads_inference():
-    """
-    Test that run_cspca.py runs without crashing using an existing YAML config.
-    """
 
-    # Path to your run_pirads.py script
-    repo_root = Path(__file__).parent.parent
-    script_path = repo_root / "run_pirads.py"
+def test_shuffling_consistency(mock_args):
+    # Verify that the image and label are shuffled with the SAME permutation
+    num_patches = 10
 
-    # Path to your existing config.yaml
-    config_path = repo_root / "config" / "config_pirads_test.yaml"  # adjust this path
+    # Distinct data per patch: [0, 1, 2, 3...]
+    data = torch.arange(num_patches).view(1, num_patches, 1, 1, 1).float()
+    target = torch.tensor([3.0])
 
-    # Make sure the file exists
-    assert config_path.exists(), f"Config file not found: {config_path}"
+    # Heatmap matches the data indices so we can track the "label"
+    heatmap = torch.arange(num_patches).view(1, num_patches, 1, 1, 1).float()
 
-    # Run the script with the config
-    result = subprocess.run(
-        [
-            sys.executable,
-            str(script_path),
-            "--mode",
-            "test",
-            "--config",
-            str(config_path),
-            "--dry_run",
-        ],
-        capture_output=True,
-        text=True,
-    )
+    att_labels, shuffled_images = get_attention_scores(data, target, heatmap, mock_args)
 
-    # Check that it ran without errors
-    assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
+    idx = (shuffled_images[0, :, 0, 0, 0] == 9.0).nonzero(as_tuple=True)[0]
+    # The attention score at that same index should be the maximum
+    assert att_labels[0, idx] == att_labels[0].max()
 
+    idx = (shuffled_images[0, :, 0, 0, 0] == 0.0).nonzero(as_tuple=True)[0]
+    # The attention score at that same index should be the minimum
+    assert att_labels[0, idx] == att_labels[0].min()
 
-def test_run_cspca_training():
+    shuffled_images = shuffled_images.cpu().squeeze()  # Shape [10]
+    att_labels = att_labels.cpu().squeeze()  # Shape [10]
+
+    sorted_vals, original_indices = torch.sort(shuffled_images)
+    sorted_labels = att_labels[original_indices]
+
+    for i in range(len(sorted_labels) - 1):
+        assert sorted_labels[i] <= sorted_labels[i + 1], (
+            f"Alignment broken at index {i}: Image val {sorted_vals[i]} has higher label than {sorted_vals[i + 1]}"
+        )
+
+
+def test_normalize_intensity_custom_masked_stats():
     """
-    Test that run_cspca.py runs without crashing using an existing YAML config.
+    Test that statistics (mean/std) are calculated ONLY from the masked region,
+    but applied to the whole image.
     """
 
-    # Path to your run_cspca.py script
-    repo_root = Path(__file__).parent.parent
-    script_path = repo_root / "run_cspca.py"
+    img = torch.zeros((2, 4, 4), dtype=torch.float32)
+    mask = torch.zeros((1, 4, 4), dtype=torch.float32)
 
-    # Path to your existing config.yaml
-    config_path = repo_root / "config" / "config_cspca_train.yaml"  # adjust this path
+    img[0, :, :] = 100.0
+    img[0, 0, 0] = 10.0
+    img[0, 0, 1] = 20.0
 
-    # Make sure the file exists
-    assert config_path.exists(), f"Config file not found: {config_path}"
+    img[1, :, :] = 50.0
+    img[1, 0, 0] = 2.0
+    img[1, 0, 1] = 4.0
 
-    # Run the script with the config
-    result = subprocess.run(
-        [
-            sys.executable,
-            str(script_path),
-            "--mode",
-            "train",
-            "--config",
-            str(config_path),
-            "--dry_run",
-        ],
-        capture_output=True,
-        text=True,
-    )
+    mask[0, 0, 0] = 1
+    mask[0, 0, 1] = 1
 
-    # Check that it ran without errors
-    assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
+    normalizer = NormalizeIntensity_custom(nonzero=False, channel_wise=True)
+    out = normalizer(img, mask)
 
+    assert torch.isclose(out[0, 0, 0], torch.tensor(-1.0)), "Ch0 masked value 1 incorrect"
+    assert torch.isclose(out[0, 0, 1], torch.tensor(1.0)), "Ch0 masked value 2 incorrect"
+    assert torch.isclose(out[0, 1, 1], torch.tensor(17.0)), "Ch0 background normalization incorrect"
 
-def test_run_cspca_inference():
+    assert torch.isclose(out[1, 0, 0], torch.tensor(-1.0)), "Ch1 masked value 1 incorrect"
+    assert torch.isclose(out[1, 1, 1], torch.tensor(47.0)), "Ch1 background normalization incorrect"
+
+
+def test_normalize_intensity_constant_area():
     """
-    Test that run_cspca.py runs without crashing using an existing YAML config.
+    Test edge case where the area under the mask has 0 variance (constant value).
+    Std should default to 1.0 to avoid division by zero.
     """
-
-    # Path to your run_cspca.py script
-    repo_root = Path(__file__).parent.parent
-    script_path = repo_root / "run_cspca.py"
-
-    # Path to your existing config.yaml
-    config_path = repo_root / "config" / "config_cspca_test.yaml"  # adjust this path
-
-    # Make sure the file exists
-    assert config_path.exists(), f"Config file not found: {config_path}"
-
-    # Run the script with the config
-    result = subprocess.run(
-        [
-            sys.executable,
-            str(script_path),
-            "--mode",
-            "test",
-            "--config",
-            str(config_path),
-            "--dry_run",
-        ],
-        capture_output=True,
-        text=True,
-    )
-
-    # Check that it ran without errors
-    assert result.returncode == 0, f"Script failed with:\n{result.stderr}"
+    img = torch.ones((1, 4, 4)) * 10.0  # All values are 10
+    mask = torch.ones((1, 4, 4))
+
+    normalizer = NormalizeIntensity_custom(channel_wise=True)
+    out = normalizer(img, mask)
+    assert torch.allclose(out, torch.zeros_like(out))
+
+    data = torch.rand(1, 10, 10)
+    mask = torch.randint(0, 2, (1, 10, 10)).float()
+    normalizer = NormalizeIntensity_custom(nonzero=False, channel_wise=True)
+    out = normalizer(data, mask)
+
+    masked = data[mask != 0]
+    mean_val = torch.mean(masked.float())
+    std_val = torch.std(masked.float(), unbiased=False)
+
+    epsilon = 1e-8
+    normalized_data = (data - mean_val) / (std_val + epsilon)
+
+    torch.testing.assert_close(out, normalized_data)
+
+
+def test_run_models():
+    args = argparse.Namespace()
+    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args.epochs = 1
+    args.batch_size = 2
+    args.tile_size = 10
+    args.tile_count = 5
+    args.use_heatmap = True
+    args.amp = False
+    args.num_classes = 4
+    args.dry_run = True
+    args.depth = 3
+
+    model = MILModel3D(num_classes=args.num_classes, mil_mode="att_trans")
+    model.to(args.device)
+    params = model.parameters()
+    loader = get_dataloader(args, split="train")
+    optimizer = torch.optim.AdamW(params, lr=1e-5, weight_decay=1e-5)
+    scaler = torch.amp.GradScaler(device=str(args.device), enabled=args.amp)
+
+    _ = train_pirads.train_epoch(model, loader, optimizer, scaler=scaler, epoch=0, args=args)
+    _ = train_pirads.val_epoch(model, loader, epoch=0, args=args)
+
+    cspca_model = CSPCAModel(backbone=model).to(args.device)
+    optimizer_cspca = torch.optim.AdamW(cspca_model.parameters(), lr=1e-5)
+    _ = train_cspca.train_epoch(cspca_model, loader, optimizer_cspca, epoch=0, args=args)
+    _ = train_cspca.val_epoch(cspca_model, loader, epoch=0, args=args)