Update README.md

README.md

@@ -28,6 +28,7 @@ By processing views sequentially with action conditioning, seq-JEPA naturally se
 | Checkpoint | Dataset | Training | Download |
 |------------|---------|----------|----------|
 | `3diebench_rot_seqlen3.pth` | 3DIEBench | seq-len=3, rotation conditioning | [Download](https://huggingface.co/Hafez/seq-JEPA/resolve/main/3diebench_rot_seqlen3.pth) |
+| `3diebench_rotcol_seqlen4.pth` | 3DIEBench | seq-len=4, rotation and color conditioning | [Download](https://huggingface.co/Hafez/seq-JEPA/resolve/main/3diebench_rotcol_seqlen4.pth) |
 | `stl10_pls.pth` | STL10 | PLS (predictive learning across saccades) | [Download](https://huggingface.co/Hafez/seq-JEPA/resolve/main/stl10_pls.pth) |
 
 ## Usage
@@ -40,15 +41,21 @@ cd seq-jepaThen load the checkpoints:
 import torch
 from models import SeqJEPA_Transforms, SeqJEPA_PLS
 
-# 3DIEBench checkpoint
+# 3DIEBench checkpoints
 kwargs = {
     "num_heads": 4, "n_channels": 3, "num_enc_layers": 3,
     "num_classes": 55, "act_cond": True, "pred_hidden": 1024,
-    "act_projdim": 128, "act_latentdim": 4, "cifar_resnet": False,
+    "act_projdim": 128, "cifar_resnet": False,
     "learn_act_emb": True
 }
+
+### for ckpt with rotation and color conditioning
+kwargs["act_latentdim"]=6
+### for ckpt with rotation conditioning
+kwargs["act_latentdim"]=4
+
 model = SeqJEPA_Transforms(img_size=128, ema=True, ema_decay=0.996, **kwargs)
-ckpt = torch.load('3diebench_rot_seqlen3.pth')
+ckpt = torch.load('3diebench_rot_seqlen3.pth') ## or ckpt = torch.load('3diebench_rotcol_seqlen4.pth') for ckpt w/ rotcolor conditioning
 model.load_state_dict(ckpt['model_state_dict'])
 
 # STL10 PLS checkpoint