Upload checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins

checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/wandb/offline-run-20260119_052527-checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins-run0/files/output.log

@@ -1,3 +1,189 @@
+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_jigsaw_swap_celoss/vlm_gym_jigsaw_swap_train
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step0
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 1.0355833768844604, mse_avg: 0.09119202941656113
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step500
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.08522484451532364, mse_avg: 0.07227052003145218
 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/
@@ -864,192 +1050,21 @@ wandb: For more information, check out the docs at: https://weave-docs.wandb.ai/
 [2026-01-19 08:27:49] (step=0000853) Train Loss mse: 0.0668, Train Loss ce: 0.0769, Train Steps/Sec: 0.08,
 [2026-01-19 08:27:59] (step=0000854) Train Loss mse: 0.0867, Train Loss ce: 0.0813, Train Steps/Sec: 0.10,
 [2026-01-19 08:28:12] (step=0000855) Train Loss mse: 0.0717, Train Loss ce: 0.0791, Train Steps/Sec: 0.07,
-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_jigsaw_swap_celoss/vlm_gym_jigsaw_swap_train
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step0
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 1.0355833768844604, mse_avg: 0.09119202941656113
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step500
+[2026-01-19 08:28:24] (step=0000856) Train Loss mse: 0.0745, Train Loss ce: 0.0779, Train Steps/Sec: 0.09,
+[2026-01-19 08:28:35] (step=0000857) Train Loss mse: 0.0827, Train Loss ce: 0.0701, Train Steps/Sec: 0.09,
+[2026-01-19 08:28:48] (step=0000858) Train Loss mse: 0.0707, Train Loss ce: 0.0751, Train Steps/Sec: 0.08,
+[2026-01-19 08:29:03] (step=0000859) Train Loss mse: 0.0528, Train Loss ce: 0.0762, Train Steps/Sec: 0.07,
+[2026-01-19 08:29:14] (step=0000860) Train Loss mse: 0.0871, Train Loss ce: 0.0803, Train Steps/Sec: 0.09,
+[2026-01-19 08:29:25] (step=0000861) Train Loss mse: 0.0612, Train Loss ce: 0.0715, Train Steps/Sec: 0.09,
+[2026-01-19 08:29:36] (step=0000862) Train Loss mse: 0.0526, Train Loss ce: 0.0828, Train Steps/Sec: 0.09,
+[2026-01-19 08:29:47] (step=0000863) Train Loss mse: 0.1003, Train Loss ce: 0.0734, Train Steps/Sec: 0.09,
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step1000
 Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
 [eval debug] first 3 batch fingerprints:
   fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
   fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
   fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 0.08522484451532364, mse_avg: 0.07227052003145218
+ce_avg: 0.11199185252189636, mse_avg: 0.06597896665334702
 base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step1500
 Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
 [eval debug] first 3 batch fingerprints:
@@ -1057,21 +1072,6 @@ Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
   fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
   fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
 ce_avg: 0.14740489423274994, mse_avg: 0.06398878246545792
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step2000
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 0.17242415249347687, mse_avg: 0.06874960660934448
-[2026-01-19 08:28:24] (step=0000856) Train Loss mse: 0.0745, Train Loss ce: 0.0779, Train Steps/Sec: 0.09,
-[2026-01-19 08:28:35] (step=0000857) Train Loss mse: 0.0827, Train Loss ce: 0.0701, Train Steps/Sec: 0.09,
-[2026-01-19 08:28:48] (step=0000858) Train Loss mse: 0.0707, Train Loss ce: 0.0751, Train Steps/Sec: 0.08,
-[2026-01-19 08:29:03] (step=0000859) Train Loss mse: 0.0528, Train Loss ce: 0.0762, Train Steps/Sec: 0.07,
-[2026-01-19 08:29:14] (step=0000860) Train Loss mse: 0.0871, Train Loss ce: 0.0803, Train Steps/Sec: 0.09,
-[2026-01-19 08:29:25] (step=0000861) Train Loss mse: 0.0612, Train Loss ce: 0.0715, Train Steps/Sec: 0.09,
-[2026-01-19 08:29:36] (step=0000862) Train Loss mse: 0.0526, Train Loss ce: 0.0828, Train Steps/Sec: 0.09,
-[2026-01-19 08:29:47] (step=0000863) Train Loss mse: 0.1003, Train Loss ce: 0.0734, Train Steps/Sec: 0.09,
 [2026-01-19 08:29:58] (step=0000864) Train Loss mse: 0.0765, Train Loss ce: 0.0725, Train Steps/Sec: 0.09,
 [2026-01-19 08:30:13] (step=0000865) Train Loss mse: 0.0489, Train Loss ce: 0.0778, Train Steps/Sec: 0.07,
 [2026-01-19 08:30:23] (step=0000866) Train Loss mse: 0.0780, Train Loss ce: 0.0754, Train Steps/Sec: 0.09,
@@ -2296,6 +2296,20 @@ ce_avg: 0.17242415249347687, mse_avg: 0.06874960660934448
 [2026-01-19 12:42:08] (step=0002085) Train Loss mse: 0.0849, Train Loss ce: 0.0714, Train Steps/Sec: 0.10,
 [2026-01-19 12:42:18] (step=0002086) Train Loss mse: 0.0803, Train Loss ce: 0.0749, Train Steps/Sec: 0.10,
 [2026-01-19 12:42:32] (step=0002087) Train Loss mse: 0.0616, Train Loss ce: 0.0706, Train Steps/Sec: 0.07,
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step2000
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.17242415249347687, mse_avg: 0.06874960660934448
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step2500
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.19842223823070526, mse_avg: 0.06336650252342224
 [2026-01-19 12:42:44] (step=0002088) Train Loss mse: 0.0685, Train Loss ce: 0.0754, Train Steps/Sec: 0.08,
 [2026-01-19 12:42:56] (step=0002089) Train Loss mse: 0.0967, Train Loss ce: 0.0695, Train Steps/Sec: 0.09,
 [2026-01-19 12:43:07] (step=0002090) Train Loss mse: 0.0986, Train Loss ce: 0.0797, Train Steps/Sec: 0.09,
@@ -2408,20 +2422,6 @@ ce_avg: 0.17242415249347687, mse_avg: 0.06874960660934448
 [2026-01-19 13:05:10] (step=0002197) Train Loss mse: 0.0664, Train Loss ce: 0.0750, Train Steps/Sec: 0.09,
 [2026-01-19 13:05:20] (step=0002198) Train Loss mse: 0.0880, Train Loss ce: 0.0704, Train Steps/Sec: 0.09,
 [2026-01-19 13:05:33] (step=0002199) Train Loss mse: 0.0675, Train Loss ce: 0.0757, Train Steps/Sec: 0.08,
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step2500
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 0.19842223823070526, mse_avg: 0.06336650252342224
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step3000
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 0.07113409787416458, mse_avg: 0.07119625806808472
 [2026-01-19 13:05:44] (step=0002200) Train Loss mse: 0.0508, Train Loss ce: 0.0733, Train Steps/Sec: 0.09,
 [2026-01-19 13:05:55] (step=0002201) Train Loss mse: 0.1194, Train Loss ce: 0.0703, Train Steps/Sec: 0.09,
 [2026-01-19 13:06:08] (step=0002202) Train Loss mse: 0.0828, Train Loss ce: 0.0721, Train Steps/Sec: 0.08,
@@ -3187,6 +3187,20 @@ ce_avg: 0.07113409787416458, mse_avg: 0.07119625806808472
 [2026-01-19 15:43:48] (step=0002959) Train Loss mse: 0.0704, Train Loss ce: 0.0697, Train Steps/Sec: 0.07,
 [2026-01-19 15:44:02] (step=0002960) Train Loss mse: 0.0800, Train Loss ce: 0.0681, Train Steps/Sec: 0.07,
 [2026-01-19 15:44:14] (step=0002961) Train Loss mse: 0.0970, Train Loss ce: 0.0651, Train Steps/Sec: 0.09,
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step3000
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.07113409787416458, mse_avg: 0.07119625806808472
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step3500
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.07160378992557526, mse_avg: 0.0721038281917572
 [2026-01-19 15:44:23] (step=0002962) Train Loss mse: 0.0561, Train Loss ce: 0.0671, Train Steps/Sec: 0.10,
 [2026-01-19 15:44:36] (step=0002963) Train Loss mse: 0.0818, Train Loss ce: 0.0704, Train Steps/Sec: 0.08,
 [2026-01-19 15:44:44] (step=0002964) Train Loss mse: 0.0788, Train Loss ce: 0.0692, Train Steps/Sec: 0.11,
@@ -3475,20 +3489,6 @@ ce_avg: 0.07113409787416458, mse_avg: 0.07119625806808472
 [2026-01-19 16:43:08] (step=0003247) Train Loss mse: 0.0538, Train Loss ce: 0.0729, Train Steps/Sec: 0.09,
 [2026-01-19 16:43:18] (step=0003248) Train Loss mse: 0.1023, Train Loss ce: 0.0670, Train Steps/Sec: 0.10,
 [2026-01-19 16:43:30] (step=0003249) Train Loss mse: 0.0763, Train Loss ce: 0.0710, Train Steps/Sec: 0.08,
-[2026-01-19 16:43:41
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step3500
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 0.07160378992557526, mse_avg: 0.0721038281917572
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step4000
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
 [2026-01-19 16:43:41] (step=0003250) Train Loss mse: 0.0678, Train Loss ce: 0.0710, Train Steps/Sec: 0.09,
 [2026-01-19 16:43:54] (step=0003251) Train Loss mse: 0.1078, Train Loss ce: 0.0676, Train Steps/Sec: 0.08,
 [2026-01-19 16:44:07] (step=0003252) Train Loss mse: 0.0562, Train Loss ce: 0.0690, Train Steps/Sec: 0.07,
@@ -4369,6 +4369,20 @@ Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
 [2026-01-19 19:45:26] (step=0004127) Train Loss mse: 0.0530, Train Loss ce: 0.0686, Train Steps/Sec: 0.09,
 [2026-01-19 19:45:37] (step=0004128) Train Loss mse: 0.0773, Train Loss ce: 0.0697, Train Steps/Sec: 0.09,
 [2026-01-19 19:45:49] (step=0004129) Train Loss mse: 0.0693, Train Loss ce: 0.0689, Train Steps/Sec: 0.08,
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step4000
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.07089916616678238, mse_avg: 0.0642944946885109
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step4500
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.07053616642951965, mse_avg: 0.0635463297367096
 [2026-01-19 19:45:59] (step=0004130) Train Loss mse: 0.0523, Train Loss ce: 0.0727, Train Steps/Sec: 0.10,
 [2026-01-19 19:46:12] (step=0004131) Train Loss mse: 0.0772, Train Loss ce: 0.0679, Train Steps/Sec: 0.07,
 [2026-01-19 19:46:26] (step=0004132) Train Loss mse: 0.0684, Train Loss ce: 0.0704, Train Steps/Sec: 0.07,
@@ -4719,20 +4733,6 @@ Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
 [2026-01-19 20:56:28] (step=0004477) Train Loss mse: 0.0562, Train Loss ce: 0.0708, Train Steps/Sec: 0.07,
 [2026-01-19 20:56:42] (step=0004478) Train Loss mse: 0.0685, Train Loss ce: 0.0723, Train Steps/Sec: 0.07,
 [2026-01-19 20:56:54] (step=0004479) Train Loss mse: 0.0507, Train Loss ce: 0.0680, Train Steps/Sec: 0.08,
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step4500
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 0.07053616642951965, mse_avg: 0.0635463297367096
-base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step5000
-Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
-[eval debug] first 3 batch fingerprints:
-  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
-ce_avg: 0.07069549709558487, mse_avg: 0.06882398575544357
 [2026-01-19 20:57:09] (step=0004480) Train Loss mse: 0.0612, Train Loss ce: 0.0761, Train Steps/Sec: 0.07,
 [2026-01-19 20:57:20] (step=0004481) Train Loss mse: 0.0472, Train Loss ce: 0.0673, Train Steps/Sec: 0.09,
 [2026-01-19 20:57:34] (step=0004482) Train Loss mse: 0.0396, Train Loss ce: 0.0751, Train Steps/Sec: 0.07,
@@ -5255,4 +5255,11 @@ ce_avg: 0.07069549709558487, mse_avg: 0.06882398575544357
 [2026-01-19 22:44:32] (step=0004999) Train Loss mse: 0.0512, Train Loss ce: 0.0650, Train Steps/Sec: 0.07,
 [2026-01-19 22:45:38] (step=0005000) Train Loss mse: 0.0639, Train Loss ce: 0.0726, Train Steps/Sec: 0.02,
 [2026-01-19 22:45:38] Saving checkpoint to /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/0005000.
-[2026-01-19 22:48:11] Done!
+[2026-01-19 22:48:11] Done!
+base_dir is /dev/shm/models/checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins/eval_used_rows, step_tag is checkpoints_vlm_gym_jigsaw_one_image_lr2e_5_ce_ins_step5000
+Preparing Dataset vlm_gym_jigsaw_swap_celoss_evalonce/vlm_gym_jigsaw_swap_val
+[eval debug] first 3 batch fingerprints:
+  fp[0]: [{'data_indexes': [0], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[1]: [{'data_indexes': [8], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+  fp[2]: [{'data_indexes': [16], 'worker_id': 0, 'dataset_name': 'vlm_gym_jigsaw_swap_celoss_evalonce'}]
+ce_avg: 0.07069549709558487, mse_avg: 0.06882398575544357