Training modifications.

.gitignore

@@ -2,3 +2,4 @@
 __pycache__/
 .DS_Store
 outputs/
+ai/

model/training.py

@@ -102,20 +102,21 @@ def test_model(
 
     test_progress = progress.get_sub_tracker("Testing Batches")
     test_progress.update(total=len(test_loader), advance=0)
-    for _, (mri, xls, targets, _) in enumerate(test_loader):
-        _check_control_events(stop_event=stop_event, pause_event=pause_event)
-        mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
-        outputs = model((mri, xls))
-        loss = criterion(outputs, targets)
-        batch_size = mri.size(0)
-        test_loss += loss.item() * batch_size
-        total_samples += batch_size
+    with torch.no_grad():
+        for _, (mri, xls, targets, _) in enumerate(test_loader):
+            _check_control_events(stop_event=stop_event, pause_event=pause_event)
+            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
+            outputs = model((mri, xls))
+            loss = criterion(outputs, targets)
+            batch_size = mri.size(0)
+            test_loss += loss.item() * batch_size
+            total_samples += batch_size
 
-        # Calculate accuracy
-        batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
-        correct += batch_correct
-        total += batch_total
-        test_progress.update(total=len(test_loader), advance=1)
+            # Calculate accuracy
+            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
+            correct += batch_correct
+            total += batch_total
+            test_progress.update(total=len(test_loader), advance=1)
 
     test_loss = test_loss / total_samples if total_samples > 0 else 0.0
     test_acc = correct / total if total > 0 else 0.0

tasks/train_normal.py

@@ -59,7 +59,6 @@ def train_normal_task(
         intermediate_model_dir.mkdir(parents=True, exist_ok=True)
     log.info(f"Intermediate models will be saved to {intermediate_model_dir}")
 
-    models = []
     train_progress = track.get_sub_tracker("Training Progress")
     track.update(total=config["training"]["ensemble_size"], advance=0)
     for model_num in range(config["training"]["ensemble_size"]):
@@ -79,8 +78,6 @@ def train_normal_task(
             .to(config["training"]["device"])
         )
 
-        models.append(model)
-
         optimizer = optim.Adam(
             model.parameters(), lr=config["training"]["learning_rate"]
         )