Update train_ms.py
This commit is contained in:
@@ -42,7 +42,7 @@ torch.backends.cudnn.allow_tf32 = True # If encontered training problem,please
|
|||||||
torch.set_float32_matmul_precision('medium')
|
torch.set_float32_matmul_precision('medium')
|
||||||
torch.backends.cuda.sdp_kernel("flash")
|
torch.backends.cuda.sdp_kernel("flash")
|
||||||
torch.backends.cuda.enable_flash_sdp(True)
|
torch.backends.cuda.enable_flash_sdp(True)
|
||||||
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
torch.backends.cuda.enable_mem_efficient_sdp(True) # Not avaliable if torch version is lower than 2.0
|
||||||
torch.backends.cuda.enable_math_sdp(True)
|
torch.backends.cuda.enable_math_sdp(True)
|
||||||
global_step = 0
|
global_step = 0
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user