Skip to content

Commit 681c000

Browse files
Merge pull request #4407 from yoinked-h/patch-1
Fix issue with 16xx cards
2 parents 37fc1fa + 62e9fec commit 681c000

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

modules/devices.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,10 +53,17 @@ def torch_gc():
5353

5454
def enable_tf32():
5555
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
5662
torch.backends.cuda.matmul.allow_tf32 = True
5763
torch.backends.cudnn.allow_tf32 = True
5864

5965

66+
6067
errors.run(enable_tf32, "Enabling TF32")
6168

6269
cpu = torch.device("cpu")

0 commit comments

Comments
 (0)