mirrors
/
SimCLR
mirror of https://github.com/sthalles/SimCLR.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix loss function labels

pull/25/head
Thalles Silva 2021-01-21 06:32:50 -03:00
parent 13a7e646e8
commit 55910e6107
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
run.py
View File

@ -54,7 +54,7 @@ parser.add_argument('--gpu-index', default=0, type=int, help='Gpu index.')
def main(): def main():
args = parser.parse_args() args = parser.parse_args()
assert args.n_views == 2, "Only two view training is supported." assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2."
# check if gpu training is available # check if gpu training is available
if not args.disable_cuda and torch.cuda.is_available(): if not args.disable_cuda and torch.cuda.is_available():
args.device = torch.device('cuda') args.device = torch.device('cuda')