From e52a1e1bd66b3a20804948a646c9e8598d411dfe Mon Sep 17 00:00:00 2001
From: Dominik Roth <dominik.roth.dev@gmail.com>
Date: Fri, 15 Oct 2021 13:34:38 +0200
Subject: [PATCH] Killed some bugs and stuff

---
 train.py | 2 +-
 utils.py | 5 ++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/train.py b/train.py
index 5c41e90..fc6c9a1 100644
--- a/train.py
+++ b/train.py
@@ -306,7 +306,7 @@ def main(argv=None):
     }
 
     # Model preparation
-    module = PetriDashNet(
+    module = PetriDishNet(
         n_in=args.n_elems,
         n_neurons=args.n_hidden,
         n_out=1,
diff --git a/utils.py b/utils.py
index e4fdf41..27f11da 100644
--- a/utils.py
+++ b/utils.py
@@ -54,7 +54,7 @@ class ParityDataset(Dataset):
 
         y = (x == 1.0).sum() % 2
 
-        return x, y
+        return x, torch.stack([y])
 
 
 class PetriDishNet(nn.Module):
@@ -110,8 +110,7 @@ class PetriDishNet(nn.Module):
                                      alpha=0.1,
                                      max_size=1e8)
 
-        self.ionDucts = ActivationSparsity(n_neurons,
-                                           alpha=0.1,
+        self.ionDucts = ActivationSparsity(alpha=0.1,
                                            beta=1.5,
                                            act_sparsity=0.65)