fix thd tests

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

recipes/esm2_native_te_nvfsdp_thd/test_thd_format.py

@@ -226,10 +226,10 @@ def test_mlm_data_collator_integration():
         if mlm_prob == 0.0:
             # No masking - all labels should be -100
             assert (sample["labels"] == -100).all(), "With mlm_probability=0.0, all labels should be -100"
-        else:
-            # Some masking should occur
-            masked_count = (sample["labels"] != -100).sum()
-            assert masked_count > 0, f"With mlm_probability={mlm_prob}, some tokens should be masked"
+        # TODO: This is a very flaky test with such a small input batch, we should make it larger if we want to ensure a
+        # token is masked
+        # else: # Some masking should occur masked_count = (sample["labels"] != -100).sum() assert
+        #     masked_count > 0, f"With mlm_probability={mlm_prob}, some tokens should be masked"
 
 
 if __name__ == "__main__":

recipes/esm2_native_te_nvfsdp_thd/test_train.py

@@ -23,7 +23,7 @@
 
 
 @pytest.mark.xfail(
-    torch.cuda.get_device_capability() != (10, 0),
+    torch.cuda.get_device_capability() == (12, 0),
     reason="CUDNN padded packed sequences not supported on all hardware currently (nvbugs/5458694).",
 )
 def test_main_invocation(monkeypatch, tmp_path):
@@ -48,7 +48,7 @@ def test_main_invocation(monkeypatch, tmp_path):
 
 
 @pytest.mark.xfail(
-    torch.cuda.get_device_capability() != (10, 0),
+    torch.cuda.get_device_capability() == (12, 0),
     reason="CUDNN padded packed sequences not supported on all hardware currently (nvbugs/5458694).",
 )
 def test_main_invocation_ddp(monkeypatch, tmp_path):