From 6bcf6b8306ce4903734fb31824799a50281cea69 Mon Sep 17 00:00:00 2001
From: fredeee
Date: Sat, 23 Mar 2024 13:27:00 +0100
Subject: add bouncingball experiment and ablation studies

---
 model/nn/eprop_transformer_shared.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

(limited to 'model/nn/eprop_transformer_shared.py')

diff --git a/model/nn/eprop_transformer_shared.py b/model/nn/eprop_transformer_shared.py
index 23a1b0f..79223c1 100644
--- a/model/nn/eprop_transformer_shared.py
+++ b/model/nn/eprop_transformer_shared.py
@@ -39,11 +39,12 @@ class EpropGateL0rdTransformerShared(nn.Module):
         self.output_embeding = OutputEmbeding(num_hidden, num_outputs) 
 
     def get_openings(self):
-        openings = 0
+        openings = []
         for i in range(self.depth):
-            openings += self.l0rds[i].l0rd.openings.item()
+            openings.append(self.l0rds[i].l0rd.openings_perslot)
 
-        return openings / self.depth
+        openings = th.mean(th.stack(openings, dim=0), dim=0)
+        return openings
 
     def get_hidden(self):
         return self.hidden
-- 
cgit v1.2.3