fix extract feature in MPS device

2023-06-24 16:21:31 +08:00 · 2023-06-24 16:21:31 +08:00 · 46c0e9b2fe
parent 359ba54321
commit 46c0e9b2fe
1 changed files with 2 additions and 0 deletions
--- a/extract_feature_print.py
+++ b/extract_feature_print.py
@ -1,5 +1,7 @@
 import os, sys, traceback

+os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
+
 # device=sys.argv[1]
 n_part = int(sys.argv[2])
 i_part = int(sys.argv[3])