lightning logs delete + predict fix

Author: mmaecki

{{cookiecutter.project_slug}}/.gitignore

@@ -104,4 +104,3 @@ instance/
 *.png
 art_checkpoints/
 lightning_logs/
-

{{cookiecutter.project_slug}}/EDA.ipynb

@@ -129,6 +129,19 @@
     "project.run_all()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from art.steps import OverfitOneBatch\n",
+    "from art.checks import CheckScoreLessThan\n",
+    "project.add_step(OverfitOneBatch(ResNet18, number_of_steps=40),\n",
+    "            [CheckScoreLessThan(metric=ce_loss, value=0.05)])\n",
+    "project.run_all()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

{{cookiecutter.project_slug}}/models/ResNet.py

@@ -27,23 +27,22 @@ def __init__(self, num_classes=100, lr=1e-3):
             transforms.Resize(256),
             transforms.CenterCrop(224),
         ])
+        # for name, para in self.model.named_parameters():
+        #     para.requires_grad = True
 
     def parse_data(self, data):
         """This is first step of your pipeline it always has batch keys inside"""
         X = data[BATCH][INPUT]
         X = X / 255
         X = rearrange(X, "b h w c -> b c h w")
         X = self.preprocess(X)
-
-        return {INPUT: X, TARGET: data[BATCH][TARGET]}
+        target = data[BATCH][TARGET].long()
+        return {INPUT: X, TARGET: target}
 
 
 
-    def predict(self, data: Dict):
-        # preds = self.model(data[INPUT]).detach().numpy()
-        preds = self.model(data[INPUT]).detach()
-        # perhaps softmax will be needed
-        return {PREDICTION: preds, TARGET: data[TARGET]}
+    def predict(self, data: Dict):       
+        return {PREDICTION: self.model(data[INPUT]), TARGET: data[TARGET]}
 
     def compute_loss(self, data):
         # Notice that the loss calculation is done in MetricsCalculator!