From c94ce8bee0b9de209c479a82e467973a16ad7708 Mon Sep 17 00:00:00 2001 From: mhz Date: Fri, 23 Aug 2024 14:11:57 +0200 Subject: [PATCH] change the device to the cuda --- correlation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/correlation.py b/correlation.py index 10a3206..c758284 100644 --- a/correlation.py +++ b/correlation.py @@ -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')