Merge pull request #13 from PolymathicAI/tutorial

fixed tutorial
PolymathicAI · Nov 2, 2023 · 843a09f · 843a09f
2 parents cb92c5c + bc45957
commit 843a09f
Show file tree

Hide file tree

Showing 2 changed files with 256 additions and 307 deletions.
diff --git a/notebooks/tutorial/AstroCLIPTutorial.ipynb b/notebooks/tutorial/AstroCLIPTutorial.ipynb
diff --git a/notebooks/tutorial/tutorial_helpers.py b/notebooks/tutorial/tutorial_helpers.py
@@ -52,11 +52,10 @@ def forward(
     pos_emb = self.position_embed(pos)  # to shape (t, embedding_dim)
 
     x = self.dropout(data_emb + pos_emb)
-    embeddings = []
     for block in self.blocks:
         x = block(x)
-        embeddings.append(x.detach().clone())
     x = self.final_layernorm(x)
+    embedding = x.detach().clone()
 
     preds = self.head(x)
     if y is not None:
@@ -66,7 +65,7 @@ def forward(
     else:
         loss = None
 
-    return {"preds": preds, "loss": loss, "embeddings": embeddings}
+    return {"preds": preds, "loss": loss, "embedding": embedding}
 
 def slice(x, section_length=10, overlap=5):
 
@@ -79,7 +78,6 @@ def slice(x, section_length=10, overlap=5):
 
     return torch.cat(sections, 1)
 
-
 def fnc(x):
     std, mean = x.std(1, keepdim=True).clip_(0.2), x.mean(1, keepdim=True)
     x = (x - mean) / std