diff --git a/tests/models/test_models_unet_2d_condition.py b/tests/models/test_models_unet_2d_condition.py
index d3ca5ea3048e..332d2f89f452 100644
--- a/tests/models/test_models_unet_2d_condition.py
+++ b/tests/models/test_models_unet_2d_condition.py
@@ -40,6 +40,7 @@
 logger = logging.get_logger(__name__)
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.use_deterministic_algorithms(True)
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = os.environ.get("CUBLAS_WORKSPACE_CONFIG", ":16:8")
 
 
 def create_lora_layers(model, mock_weights: bool = True):