Disable TF32 on DDP tests

zasdfgbnm · zasdfgbnm · commit 6c0052a1a6c0 · 2021-02-26T12:40:07.000-08:00
diff --git a/test/distributed/test_distributed_fork.py b/test/distributed/test_distributed_fork.py
@@ -11,6 +11,8 @@
     DistributedTest, TestDistBackend
 )
 
+torch.backends.cuda.matmul.allow_tf32 = False
+
 CPP_EXTENSIONS_WARNING = """
 Ninja (https://ninja-build.org) must be available to run C++ extensions tests,
 but it could not be found. Install ninja with `pip install ninja`
@@ -48,6 +50,7 @@ class TestDistBackendWithFork(TestDistBackend, DistributedTest._DistTestBase):
         def setUp(self):
             super().setUp()
             self._fork_processes()
+            torch.backends.cudnn.flags(allow_tf32=False).__enter__()
 
 
 elif BACKEND == "mpi":
diff --git a/test/distributed/test_distributed_spawn.py b/test/distributed/test_distributed_spawn.py
@@ -3,12 +3,15 @@
 import sys
 import unittest
 
+import torch
 import torch.distributed as dist
 from torch.testing._internal.common_utils import run_tests, TEST_WITH_ASAN, NO_MULTIPROCESSING_SPAWN
 from torch.testing._internal.distributed.distributed_test import (
     DistributedTest, TestDistBackend
 )
 
+torch.backends.cuda.matmul.allow_tf32 = False
+
 if not dist.is_available():
     print("Distributed not available, skipping tests", file=sys.stderr)
     sys.exit(0)
@@ -28,6 +31,7 @@ class TestDistBackendWithSpawn(TestDistBackend, DistributedTest._DistTestBase):
         def setUp(self):
             super().setUp()
             self._spawn_processes()
+            torch.backends.cudnn.flags(allow_tf32=False).__enter__()
 
 
 if __name__ == "__main__":
