From 9e5937531155ac334aab5fe3ac5490e32adb8a34 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Sat, 1 Apr 2023 17:01:51 +0800
Subject: [PATCH] fix: extract freture cannot run on pure cpu

---
 extract_feature_print.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/extract_feature_print.py b/extract_feature_print.py
index 8ee1122..a92f272 100644
--- a/extract_feature_print.py
+++ b/extract_feature_print.py
@@ -1,9 +1,14 @@
 import os,sys,traceback
-n_part=int(sys.argv[1])
-i_part=int(sys.argv[2])
-i_gpu=sys.argv[3]
-exp_dir=sys.argv[4]
-os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)
+if len(sys.argv) == 4:
+    n_part=int(sys.argv[1])
+    i_part=int(sys.argv[2])
+    exp_dir=sys.argv[3]
+else:
+    n_part=int(sys.argv[1])
+    i_part=int(sys.argv[2])
+    i_gpu=sys.argv[3]
+    exp_dir=sys.argv[4]
+    os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)
 
 import torch
 import torch.nn.functional as F