Merge pull request #28 from danny-1k/dev

IJEPA implementation
HMUNACHI · Mar 25, 2024 · 8967b1d · 8967b1d
2 parents dc47e23 + 987d438
commit 8967b1d
Show file tree

Hide file tree

Showing 3 changed files with 802 additions and 0 deletions.
diff --git a/nanodl/__init__.py b/nanodl/__init__.py
@@ -133,6 +133,12 @@
     RewardDataParallelTrainer
 )
 
+from nanodl.__src.models.ijepa import (
+    IJEPA,
+    IJEPADataParallelTrainer,
+    IJEPADataSampler
+)
+
 from nanodl.__src.layers.attention import (
     MultiQueryAttention,
     LocalMultiHeadAttention,
@@ -193,6 +199,9 @@
     "GaussianProcess",
 
     # Models
+    "IJEPA",
+    "IJEPADataParallelTrainer",
+    "IJEPADataSampler",
     "Gemma",
     "GemmaDataParallelTrainer",
     "GemmaDecoder",