Update on "[NJT] Fix inference mode for composite implicit ops without nested-specific kernel"

soulitzer · soulitzer · commit 9e27eb7b7d56 · 2025-02-06T16:02:31.000-05:00
[ghstack-poisoned]
diff --git a/test/test_nestedtensor.py b/test/test_nestedtensor.py
@@ -7628,10 +7628,14 @@ def test_dropout_inference_mode(self, device):
         seq_len = 32
         embed_dim = 128
 
-        nt = torch.nested.nested_tensor([
-            torch.randn(11, seq_len, embed_dim, device=device),
-            torch.randn(11, seq_len, embed_dim, device=device)
-        ], layout=torch.jagged, device=device)
+        nt = torch.nested.nested_tensor(
+            [
+                torch.randn(11, seq_len, embed_dim, device=device),
+                torch.randn(11, seq_len, embed_dim, device=device),
+            ],
+            layout=torch.jagged,
+            device=device,
+        )
 
         with torch.inference_mode():
             torch.nn.functional.dropout(nt, p=0.05)
