From 5914ed3efc47591d3f4cc3bf5f2777d6d9f86b48 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <2225664821@qq.com>
Date: Mon, 4 Sep 2023 14:27:14 +0800
Subject: [PATCH] speed optim

---
 train_ms.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/train_ms.py b/train_ms.py
index 55f2bb5..715f51f 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -40,6 +40,10 @@ torch.backends.cudnn.benchmark = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 torch.set_float32_matmul_precision('medium')
+torch.backends.cuda.sdp_kernel("flash")
+torch.backends.cuda.enable_flash_sdp(True)
+torch.backends.cuda.enable_mem_efficient_sdp(True)
+torch.backends.cuda.enable_math_sdp(True)
 global_step = 0