Updates to jepas compatability with python

demo_jepa_encoder.py

@@ -5,10 +5,18 @@ encoder = JepaEncoder.load_model(
 )
 
 import numpy
+import torch
 img = numpy.random.random(size=(360, 480, 3))
 
+x = torch.rand((32, 3, 256, 900))
+
 print("Input Img:", img.shape)
 embedding = encoder.embed_image(img)
 
+print(embedding)
+print(embedding.shape)
+
+
+embedding = encoder.embed_image(x)
 print(embedding)
 print(embedding.shape)

params-encoder.yaml

@@ -1,11 +1,6 @@
 app: vjepa
 data:
-  batch_size: 8
-  clip_duration: null
   crop_size: 224
-  dataset_type: VideoDataset
-  datasets:
-  - /path/to/dataset.csv
   decode_one_clip: true
   filter_short_videos: false
   num_clips: 1
@@ -14,7 +9,7 @@ data:
   patch_size: 16
   pin_mem: true
   sampling_rate: 4
-  tubelet_size: 2
+  tubelet_size: 1
 data_aug:
   auto_augment: false
   motion_shift: false
@@ -26,11 +21,8 @@ data_aug:
   - 1.0
   reprob: 0.0
 logging:
-  folder: /path/to/logs
+  folder: /media/rpal/Drive_10TB/John/jepa/logs
   write_tag: jepa
-loss:
-  loss_exp: 1.0
-  reg_coeff: 0.0
 mask:
 - aspect_ratio:
   - 0.75
@@ -60,7 +52,7 @@ meta:
   dtype: bfloat16
   eval_freq: 100
   load_checkpoint: true
-  read_checkpoint: /path/to/vitl16.pth.tar
+  read_checkpoint: /media/rpal/Drive_10TB/John/jepa/huggingface/jepa-latest.pth.tar
   save_every_freq: 5
   seed: 234
   use_sdpa: true
@@ -71,19 +63,3 @@ model:
   uniform_power: true
   use_mask_tokens: true
   zero_init_mask_tokens: true
-nodes: 16
-optimization:
-  clip_grad: 10.0
-  ema:
-  - 0.998
-  - 1.0
-  epochs: 25
-  final_lr: 1.0e-06
-  final_weight_decay: 0.4
-  ipe: 300
-  ipe_scale: 1.25
-  lr: 0.000625
-  start_lr: 0.0002
-  warmup: 40
-  weight_decay: 0.04
-tasks_per_node: 8