Upload checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins

checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/wandb/offline-run-20260129_220927-checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins-run0/files/output.log

@@ -1,189 +1,3 @@
-[2026-01-30 02:11:22] (step=0000885) Train Loss mse: 0.0253, Train Loss ce: 0.0419, Train Steps/Sec: 0.06,
-FullyShardedDataParallel(
-  (_fsdp_wrapped_module): Bagel(
-    (language_model): Qwen2ForCausalLM(
-      (model): Qwen2Model(
-        (embed_tokens): Embedding(152064, 3584)
-        (layers): ModuleList(
-          (0-27): 28 x FullyShardedDataParallel(
-            (_fsdp_wrapped_module): CheckpointWrapper(
-              (_checkpoint_wrapped_module): Qwen2MoTDecoderLayer(
-                (self_attn): PackedAttentionMoT(
-                  (q_proj): Linear(in_features=3584, out_features=3584, bias=True)
-                  (k_proj): Linear(in_features=3584, out_features=512, bias=True)
-                  (v_proj): Linear(in_features=3584, out_features=512, bias=True)
-                  (o_proj): Linear(in_features=3584, out_features=3584, bias=False)
-                  (q_norm): Qwen2RMSNorm((128,), eps=1e-06)
-                  (k_norm): Qwen2RMSNorm((128,), eps=1e-06)
-                  (q_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
-                  (k_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
-                  (q_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=True)
-                  (k_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
-                  (v_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
-                  (o_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=False)
-                )
-                (mlp): Qwen2MLP(
-                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
-                  (act_fn): SiLU()
-                )
-                (mlp_moe_gen): Qwen2MLP(
-                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
-                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
-                  (act_fn): SiLU()
-                )
-                (input_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
-                (input_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-                (post_attention_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
-                (post_attention_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-              )
-            )
-          )
-        )
-        (norm): Qwen2RMSNorm((3584,), eps=1e-06)
-        (norm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
-        (rotary_emb): Qwen2RotaryEmbedding()
-      )
-      (lm_head): Linear(in_features=3584, out_features=152064, bias=False)
-    )
-    (time_embedder): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): TimestepEmbedder(
-        (mlp): Sequential(
-          (0): Linear(in_features=256, out_features=3584, bias=True)
-          (1): SiLU()
-          (2): Linear(in_features=3584, out_features=3584, bias=True)
-        )
-      )
-    )
-    (vae2llm): Linear(in_features=64, out_features=3584, bias=True)
-    (llm2vae): Linear(in_features=3584, out_features=64, bias=True)
-    (latent_pos_embed): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): PositionEmbedding()
-    )
-    (vit_model): SiglipVisionModel(
-      (vision_model): FullyShardedDataParallel(
-        (_fsdp_wrapped_module): SiglipVisionTransformer(
-          (embeddings): SiglipVisionEmbeddings(
-            (position_embedding): Embedding(4900, 1152)
-            (patch_embedding): Linear(in_features=588, out_features=1152, bias=True)
-          )
-          (encoder): SiglipEncoder(
-            (layers): ModuleList(
-              (0-25): 26 x FullyShardedDataParallel(
-                (_fsdp_wrapped_module): CheckpointWrapper(
-                  (_checkpoint_wrapped_module): SiglipEncoderLayer(
-                    (self_attn): SiglipFlashAttention2(
-                      (k_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (v_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (q_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                      (out_proj): Linear(in_features=1152, out_features=1152, bias=True)
-                    )
-                    (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-                    (mlp): SiglipMLP(
-                      (activation_fn): PytorchGELUTanh()
-                      (fc1): Linear(in_features=1152, out_features=4304, bias=True)
-                      (fc2): Linear(in_features=4304, out_features=1152, bias=True)
-                    )
-                    (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-                  )
-                )
-              )
-            )
-          )
-          (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
-        )
-      )
-    )
-    (connector): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): CheckpointWrapper(
-        (_checkpoint_wrapped_module): MLPconnector(
-          (activation_fn): PytorchGELUTanh()
-          (fc1): Linear(in_features=1152, out_features=3584, bias=True)
-          (fc2): Linear(in_features=3584, out_features=3584, bias=True)
-        )
-      )
-    )
-    (vit_pos_embed): FullyShardedDataParallel(
-      (_fsdp_wrapped_module): PositionEmbedding()
-    )
-  )
-)
-_flat_param True
-language_model.model.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.26._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-language_model.model.layers.27._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-time_embedder._fsdp_wrapped_module._flat_param True
-latent_pos_embed._fsdp_wrapped_module._flat_param False
-vit_model.vision_model._fsdp_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_model.vision_model._fsdp_wrapped_module.encoder.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-connector._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
-vit_pos_embed._fsdp_wrapped_module._flat_param False
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss/vlm_gym_mental_rotation_3d_pad3_by_axis_train
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step0
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 1.9196269512176514, mse_avg: 1.8925566673278809
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step500
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
 wandb: Detected [huggingface_hub.inference] in use.
 wandb: Use W&B Weave for improved LLM call tracing. Install Weave with `pip install weave` then add `import weave` to the top of your script.
 wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
@@ -1033,6 +847,192 @@ wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
 [2026-01-30 01:58:31] (step=0000836) Train Loss mse: 0.0274, Train Loss ce: 0.0452, Train Steps/Sec: 0.07,
 [2026-01-30 01:58:46] (step=0000837) Train Loss mse: 0.0267, Train Loss ce: 0.0441, Train Steps/Sec: 0.07,
 [2026-01-30 01:59:02] (step=0000838) Train Loss mse: 0.0272, Train Loss ce: 0.0427, Train Steps/Sec: 0.06,
+FullyShardedDataParallel(
+  (_fsdp_wrapped_module): Bagel(
+    (language_model): Qwen2ForCausalLM(
+      (model): Qwen2Model(
+        (embed_tokens): Embedding(152064, 3584)
+        (layers): ModuleList(
+          (0-27): 28 x FullyShardedDataParallel(
+            (_fsdp_wrapped_module): CheckpointWrapper(
+              (_checkpoint_wrapped_module): Qwen2MoTDecoderLayer(
+                (self_attn): PackedAttentionMoT(
+                  (q_proj): Linear(in_features=3584, out_features=3584, bias=True)
+                  (k_proj): Linear(in_features=3584, out_features=512, bias=True)
+                  (v_proj): Linear(in_features=3584, out_features=512, bias=True)
+                  (o_proj): Linear(in_features=3584, out_features=3584, bias=False)
+                  (q_norm): Qwen2RMSNorm((128,), eps=1e-06)
+                  (k_norm): Qwen2RMSNorm((128,), eps=1e-06)
+                  (q_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
+                  (k_norm_moe_gen): Qwen2RMSNorm((128,), eps=1e-06)
+                  (q_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=True)
+                  (k_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
+                  (v_proj_moe_gen): Linear(in_features=3584, out_features=512, bias=True)
+                  (o_proj_moe_gen): Linear(in_features=3584, out_features=3584, bias=False)
+                )
+                (mlp): Qwen2MLP(
+                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
+                  (act_fn): SiLU()
+                )
+                (mlp_moe_gen): Qwen2MLP(
+                  (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
+                  (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
+                  (act_fn): SiLU()
+                )
+                (input_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
+                (input_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+                (post_attention_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
+                (post_attention_layernorm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+              )
+            )
+          )
+        )
+        (norm): Qwen2RMSNorm((3584,), eps=1e-06)
+        (norm_moe_gen): Qwen2RMSNorm((3584,), eps=1e-06)
+        (rotary_emb): Qwen2RotaryEmbedding()
+      )
+      (lm_head): Linear(in_features=3584, out_features=152064, bias=False)
+    )
+    (time_embedder): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): TimestepEmbedder(
+        (mlp): Sequential(
+          (0): Linear(in_features=256, out_features=3584, bias=True)
+          (1): SiLU()
+          (2): Linear(in_features=3584, out_features=3584, bias=True)
+        )
+      )
+    )
+    (vae2llm): Linear(in_features=64, out_features=3584, bias=True)
+    (llm2vae): Linear(in_features=3584, out_features=64, bias=True)
+    (latent_pos_embed): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): PositionEmbedding()
+    )
+    (vit_model): SiglipVisionModel(
+      (vision_model): FullyShardedDataParallel(
+        (_fsdp_wrapped_module): SiglipVisionTransformer(
+          (embeddings): SiglipVisionEmbeddings(
+            (position_embedding): Embedding(4900, 1152)
+            (patch_embedding): Linear(in_features=588, out_features=1152, bias=True)
+          )
+          (encoder): SiglipEncoder(
+            (layers): ModuleList(
+              (0-25): 26 x FullyShardedDataParallel(
+                (_fsdp_wrapped_module): CheckpointWrapper(
+                  (_checkpoint_wrapped_module): SiglipEncoderLayer(
+                    (self_attn): SiglipFlashAttention2(
+                      (k_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (v_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (q_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                      (out_proj): Linear(in_features=1152, out_features=1152, bias=True)
+                    )
+                    (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+                    (mlp): SiglipMLP(
+                      (activation_fn): PytorchGELUTanh()
+                      (fc1): Linear(in_features=1152, out_features=4304, bias=True)
+                      (fc2): Linear(in_features=4304, out_features=1152, bias=True)
+                    )
+                    (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+                  )
+                )
+              )
+            )
+          )
+          (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)
+        )
+      )
+    )
+    (connector): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): CheckpointWrapper(
+        (_checkpoint_wrapped_module): MLPconnector(
+          (activation_fn): PytorchGELUTanh()
+          (fc1): Linear(in_features=1152, out_features=3584, bias=True)
+          (fc2): Linear(in_features=3584, out_features=3584, bias=True)
+        )
+      )
+    )
+    (vit_pos_embed): FullyShardedDataParallel(
+      (_fsdp_wrapped_module): PositionEmbedding()
+    )
+  )
+)
+_flat_param True
+language_model.model.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.26._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+language_model.model.layers.27._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+time_embedder._fsdp_wrapped_module._flat_param True
+latent_pos_embed._fsdp_wrapped_module._flat_param False
+vit_model.vision_model._fsdp_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.0._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.1._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.2._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.3._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.4._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.5._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.6._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.7._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.8._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.9._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.10._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.11._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.12._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.13._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.14._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.15._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.16._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.17._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.18._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.19._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.20._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.21._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.22._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.23._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.24._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_model.vision_model._fsdp_wrapped_module.encoder.layers.25._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+connector._fsdp_wrapped_module._checkpoint_wrapped_module._flat_param True
+vit_pos_embed._fsdp_wrapped_module._flat_param False
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss/vlm_gym_mental_rotation_3d_pad3_by_axis_train
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step0
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 1.9196269512176514, mse_avg: 1.8925566673278809
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step500
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.05214770883321762, mse_avg: 0.02520925924181938
 [2026-01-30 01:59:17] (step=0000839) Train Loss mse: 0.0284, Train Loss ce: 0.0449, Train Steps/Sec: 0.07,
 [2026-01-30 01:59:33] (step=0000840) Train Loss mse: 0.0270, Train Loss ce: 0.0434, Train Steps/Sec: 0.06,
 [2026-01-30 01:59:48] (step=0000841) Train Loss mse: 0.0291, Train Loss ce: 0.0439, Train Steps/Sec: 0.07,
@@ -1079,7 +1079,7 @@ wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
 [2026-01-30 02:10:32] (step=0000882) Train Loss mse: 0.0286, Train Loss ce: 0.0426, Train Steps/Sec: 0.06,
 [2026-01-30 02:10:49] (step=0000883) Train Loss mse: 0.0296, Train Loss ce: 0.0417, Train Steps/Sec: 0.06,
 [2026-01-30 02:11:06] (step=0000884) Train Loss mse: 0.0282, Train Loss ce: 0.0432, Train Steps/Sec: 0.06,
-ce_avg: 0.05214770883321762, mse_avg: 0.02520925924181938
+[2026-01-30 02:11:22] (step=0000885) Train Loss mse: 0.0253, Train Loss ce: 0.0419, Train Steps/Sec: 0.06,
 [2026-01-30 02:11:37] (step=0000886) Train Loss mse: 0.0290, Train Loss ce: 0.0417, Train Steps/Sec: 0.07,
 [2026-01-30 02:11:54] (step=0000887) Train Loss mse: 0.0262, Train Loss ce: 0.0458, Train Steps/Sec: 0.06,
 [2026-01-30 02:12:10] (step=0000888) Train Loss mse: 0.0255, Train Loss ce: 0.0442, Train Steps/Sec: 0.06,
@@ -2202,6 +2202,41 @@ ce_avg: 0.05214770883321762, mse_avg: 0.02520925924181938
 [2026-01-30 07:10:50] (step=0002005) Train Loss mse: 0.0258, Train Loss ce: 0.0349, Train Steps/Sec: 0.06,
 [2026-01-30 07:11:04] (step=0002006) Train Loss mse: 0.0263, Train Loss ce: 0.0371, Train Steps/Sec: 0.07,
 [2026-01-30 07:11:21] (step=0002007) Train Loss mse: 0.0258, Train Loss ce: 0.0349, Train Steps/Sec: 0.06,
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step1000
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.04393863305449486, mse_avg: 0.02561650611460209
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step1500
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.041284698992967606, mse_avg: 0.02522081509232521
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step2000
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.044035617262125015, mse_avg: 0.0255599282681942
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step2500
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.04285125434398651, mse_avg: 0.02553764171898365
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step3000
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.02865723706781864, mse_avg: 0.023178618401288986
 [2026-01-30 07:11:35] (step=0002008) Train Loss mse: 0.0244, Train Loss ce: 0.0373, Train Steps/Sec: 0.07,
 [2026-01-30 07:11:51] (step=0002009) Train Loss mse: 0.0275, Train Loss ce: 0.0348, Train Steps/Sec: 0.06,
 [2026-01-30 07:12:07] (step=0002010) Train Loss mse: 0.0234, Train Loss ce: 0.0330, Train Steps/Sec: 0.06,
@@ -2288,41 +2323,6 @@ ce_avg: 0.05214770883321762, mse_avg: 0.02520925924181938
 [2026-01-30 07:33:28] (step=0002091) Train Loss mse: 0.0288, Train Loss ce: 0.0344, Train Steps/Sec: 0.07,
 [2026-01-30 07:33:43] (step=0002092) Train Loss mse: 0.0256, Train Loss ce: 0.0389, Train Steps/Sec: 0.07,
 [2026-01-30 07:33:58] (step=0002093) Train Loss mse: 0.0242, Train Loss ce: 0.0347, Train Steps/Sec: 0.07,
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step1000
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.04393863305449486, mse_avg: 0.02561650611460209
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step1500
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.041284698992967606, mse_avg: 0.02522081509232521
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step2000
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.044035617262125015, mse_avg: 0.0255599282681942
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step2500
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.04285125434398651, mse_avg: 0.02553764171898365
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step3000
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.02865723706781864, mse_avg: 0.023178618401288986
 [2026-01-30 07:34:13] (step=0002094) Train Loss mse: 0.0285, Train Loss ce: 0.0347, Train Steps/Sec: 0.06,
 [2026-01-30 07:34:29] (step=0002095) Train Loss mse: 0.0258, Train Loss ce: 0.0343, Train Steps/Sec: 0.06,
 [2026-01-30 07:34:45] (step=0002096) Train Loss mse: 0.0255, Train Loss ce: 0.0345, Train Steps/Sec: 0.06,
@@ -3282,6 +3282,20 @@ ce_avg: 0.02865723706781864, mse_avg: 0.023178618401288986
 [2026-01-30 11:50:17] (step=0003047) Train Loss mse: 0.0233, Train Loss ce: 0.0305, Train Steps/Sec: 0.06,
 [2026-01-30 11:50:34] (step=0003048) Train Loss mse: 0.0250, Train Loss ce: 0.0300, Train Steps/Sec: 0.06,
 [2026-01-30 11:50:51] (step=0003049) Train Loss mse: 0.0231, Train Loss ce: 0.0292, Train Steps/Sec: 0.06,
+[2026-01-30 11:51:08
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step3500
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.026425831019878387, mse_avg: 0.02323519065976143
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step4000
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
 [2026-01-30 11:51:08] (step=0003050) Train Loss mse: 0.0247, Train Loss ce: 0.0299, Train Steps/Sec: 0.06,
 [2026-01-30 11:51:23] (step=0003051) Train Loss mse: 0.0272, Train Loss ce: 0.0311, Train Steps/Sec: 0.07,
 [2026-01-30 11:51:39] (step=0003052) Train Loss mse: 0.0257, Train Loss ce: 0.0326, Train Steps/Sec: 0.06,
@@ -3483,20 +3497,6 @@ ce_avg: 0.02865723706781864, mse_avg: 0.023178618401288986
 [2026-01-30 12:42:51] (step=0003248) Train Loss mse: 0.0238, Train Loss ce: 0.0297, Train Steps/Sec: 0.06,
 [2026-01-30 12:43:07] (step=0003249) Train Loss mse: 0.0265, Train Loss ce: 0.0288, Train Steps/Sec: 0.06,
 [2026-01-30 12:43:22] (step=0003250) Train Loss mse: 0.0226, Train Loss ce: 0.0293, Train Steps/Sec: 0.06,
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step3500
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.026425831019878387, mse_avg: 0.02323519065976143
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step4000
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.024867402389645576, mse_avg: 0.02298266813158989
 [2026-01-30 12:43:37] (step=0003251) Train Loss mse: 0.0252, Train Loss ce: 0.0305, Train Steps/Sec: 0.07,
 [2026-01-30 12:43:53] (step=0003252) Train Loss mse: 0.0250, Train Loss ce: 0.0315, Train Steps/Sec: 0.06,
 [2026-01-30 12:44:09] (step=0003253) Train Loss mse: 0.0243, Train Loss ce: 0.0290, Train Steps/Sec: 0.06,
@@ -4344,6 +4344,20 @@ ce_avg: 0.024867402389645576, mse_avg: 0.02298266813158989
 [2026-01-30 16:28:31] (step=0004095) Train Loss mse: 0.0252, Train Loss ce: 0.0258, Train Steps/Sec: 0.07,
 [2026-01-30 16:28:47] (step=0004096) Train Loss mse: 0.0278, Train Loss ce: 0.0249, Train Steps/Sec: 0.06,
 [2026-01-30 16:29:02] (step=0004097) Train Loss mse: 0.0232, Train Loss ce: 0.0263, Train Steps/Sec: 0.07,
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step4500
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.023800626397132874, mse_avg: 0.023024622350931168
+base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step5000
+Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
+ce_avg: 0.023549562320113182, mse_avg: 0.023148616775870323
 [2026-01-30 16:29:17] (step=0004098) Train Loss mse: 0.0263, Train Loss ce: 0.0237, Train Steps/Sec: 0.07,
 [2026-01-30 16:29:32] (step=0004099) Train Loss mse: 0.0239, Train Loss ce: 0.0256, Train Steps/Sec: 0.07,
 [2026-01-30 16:29:48] (step=0004100) Train Loss mse: 0.0235, Train Loss ce: 0.0238, Train Steps/Sec: 0.06,
@@ -4494,20 +4508,6 @@ ce_avg: 0.024867402389645576, mse_avg: 0.02298266813158989
 [2026-01-30 17:07:58] (step=0004245) Train Loss mse: 0.0266, Train Loss ce: 0.0241, Train Steps/Sec: 0.07,
 [2026-01-30 17:08:14] (step=0004246) Train Loss mse: 0.0219, Train Loss ce: 0.0259, Train Steps/Sec: 0.06,
 [2026-01-30 17:08:30] (step=0004247) Train Loss mse: 0.0266, Train Loss ce: 0.0243, Train Steps/Sec: 0.06,
-[2026-01-30 17:08:46
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step4500
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-ce_avg: 0.023800626397132874, mse_avg: 0.023024622350931168
-base_dir is /dev/shm/models/checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_mental_rotation_3d_pad3_by_axis_one_image_lr2e_5_ce_ins_step5000
-Preparing Dataset vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce/vlm_gym_mental_rotation_3d_pad3_by_axis_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_mental_rotation_3d_pad3_by_axis_celoss_evalonce'}]
 [2026-01-30 17:08:46] (step=0004248) Train Loss mse: 0.0252, Train Loss ce: 0.0243, Train Steps/Sec: 0.06,
 [2026-01-30 17:09:01] (step=0004249) Train Loss mse: 0.0238, Train Loss ce: 0.0293, Train Steps/Sec: 0.07,
 [2026-01-30 17:09:16] (step=0004250) Train Loss mse: 0.0288, Train Loss ce: 0.0242, Train Steps/Sec: 0.07,