We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
2 parents 37fc1fa + 62e9fec commit 681c000Copy full SHA for 681c000
modules/devices.py
@@ -53,10 +53,17 @@ def torch_gc():
53
54
def enable_tf32():
55
if torch.cuda.is_available():
56
+ for devid in range(0,torch.cuda.device_count()):
57
+ if torch.cuda.get_device_capability(devid) == (7, 5):
58
+ shd = True
59
+ if shd:
60
+ torch.backends.cudnn.benchmark = True
61
+ torch.backends.cudnn.enabled = True
62
torch.backends.cuda.matmul.allow_tf32 = True
63
torch.backends.cudnn.allow_tf32 = True
64
65
66
+
67
errors.run(enable_tf32, "Enabling TF32")
68
69
cpu = torch.device("cpu")
0 commit comments