unit test fixes

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

Author: adil-a

nemo_automodel/components/models/biencoder/llama_bidirectional_model.py

@@ -177,7 +177,7 @@ def _update_causal_mask(
             return attention_mask
         return None
 
-    @check_model_inputs
+    @check_model_inputs()
     @auto_docstring
     def forward(
         self,

tests/unit_tests/distributed/test_utils.py

@@ -117,7 +117,7 @@ def test_get_sync_ctx(monkeypatch, patch_dist):
     class Plain(torch.nn.Linear):
         pass
 
-    ctx = du.get_sync_ctx(Plain(2, 2), is_optim_step=False)
+    ctx = du.get_sync_ctx(Plain(2, 2), is_optim_step=False, defer_fsdp_grad_sync=False)
     # entering/exiting the context must be a no-op
     with ctx:
         pass

tests/unit_tests/recipes/test_finetune_vlm_helpers.py

@@ -228,6 +228,7 @@ def test_run_train_step_supports_tensor_outputs(monkeypatch):
     recipe.cfg = _Cfg(fp8=None)
     recipe.lr_scheduler = None
     recipe.timestamp = 0.0
+    recipe.model_wrapper = None
 
     recipe._dp_allreduce = lambda tensor, include_cp=False: tensor
     recipe._get_dp_group_size = lambda include_cp=True: 1
@@ -251,7 +252,7 @@ def fake_calculate_loss(*args, **kwargs):
     )
     monkeypatch.setattr(
         "nemo_automodel.recipes.vlm.finetune.get_sync_ctx",
-        lambda model, is_last: nullcontext(),
+        lambda model, is_last, defer_fsdp_grad_sync=True: nullcontext(),
     )
 
     calculate_mock = MagicMock(side_effect=fake_calculate_loss)