From 27d8d80a4000fe4e4303fb61409c1242ac94b85d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 3 Jun 2026 02:31:11 +0000 Subject: [PATCH] Fix missing DEVICE constant in production FMHA test --- tests/unit/test_production_fmha_layer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/test_production_fmha_layer.py b/tests/unit/test_production_fmha_layer.py index c548749a..879182e6 100644 --- 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):