test_video_mae: use OH_TEST_DEVICE

Signed-off-by: U. Artie Eoff <[email protected]>
uartie · Feb 18, 2025 · 63ebaaf · 63ebaaf
1 parent 81877f4
commit 63ebaaf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_video_mae.py b/tests/test_video_mae.py
@@ -23,9 +23,9 @@
 import pytest
 import torch
 from transformers import VideoMAEForVideoClassification, VideoMAEImageProcessor
+from .utils import OH_TEST_DEVICE
 
-
-if os.environ.get("GAUDI2_CI", "0") == "1":
+if OH_TEST_DEVICE in ["gaudi2"]:
     # Gaudi2 CI baselines
     LATENCY_VIDEOMAE_BF16_GRAPH_BASELINE = 17.544198036193848
 else: