teeny adjustments

Author: hzavadil98

main.py

@@ -5,6 +5,7 @@
 from torch.utils.data import DataLoader
 from torchvision import transforms
 from tqdm import tqdm
+#from wandb_api import WANDB_API
 
 from CollaborativeCoding import (
     MetricWrapper,
@@ -121,11 +122,11 @@ def main():
             train_metrics(y, logits)
 
             break
-        print(train_metrics.accumulate())
+        print(train_metrics.__getmetrics__())
         print("Dry run completed successfully.")
         exit()
 
-    # wandb.login(key=WANDB_API)
+#    wandb.login(key=WANDB_API)
     wandb.init(
         entity="ColabCode",
         project=args.run_name,

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
     "torch>=2.6.0",
     "torchvision>=0.21.0",
     "tqdm>=4.67.1",
+    "wandb>=0.19.6",
 ]
 [tool.isort]
 profile = "black"

tests/test_models.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from CollaborativeCoding.models import ChristianModel, JanModel, MagnusModel
+from CollaborativeCoding.models import ChristianModel, JanModel, MagnusModel, SolveigModel
 
 
 @pytest.mark.parametrize(