Fix missing DEVICE constant in production FMHA test

2026-06-03 02:31:11 +00:00
parent 26a817c2f2
commit 27d8d80a40
1 changed files with 1 additions and 0 deletions
--- a/tests/unit/test_production_fmha_layer.py
+++ b/tests/unit/test_production_fmha_layer.py
@@ -20,6 +20,7 @@ import torch.nn.functional as F
 CHECKPOINT_DIR = os.environ.get(
    "CHECKPOINT_DIR", "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4")
 NUM_GPUS = int(os.environ.get("NUM_GPUS", "8"))
+DEVICE = "cuda:0"


 def cosine(a, b):