iicd/swap-nas
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

change the device to the cuda

Browse Source
This commit is contained in:
mhz 2024-08-23 14:11:57 +02:00
parent 7b2bfd05d5
commit c94ce8bee0
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
correlation.py
View File

@ -21,7 +21,7 @@ parser = argparse.ArgumentParser()
# general setting
parser.add_argument('--data_path', default="datasets", type=str, nargs='?', help='path to the image dataset (datasets or datasets/ILSVRC/Data/CLS-LOC)')
parser.add_argument('--seed', default=0, type=int, help='random seed')
parser.add_argument('--device', default="mps", type=str, nargs='?', help='setup device (cpu, mps or cuda)')
parser.add_argument('--device', default="cuda:2", type=str, nargs='?', help='setup device (cpu, mps or cuda)')
parser.add_argument('--repeats', default=32, type=int, nargs='?', help='times of calculating the training-free metric')
parser.add_argument('--input_samples', default=16, type=int, nargs='?', help='input batch size for training-free metric')