module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found one of them on device: cuda:1 about big_transfer HOT 1 OPEN

well, i find the following information:

https://stackoverflow.com/questions/59249563/runtimeerror-module-must-have-its-parameters-and-buffers-on-device-cuda1-devi

just modify the code as following:

#device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
#logger.info(f"Going to train on {device}")
device_ids = [1, 4, 5]
device = torch.device("cuda:1")
logger.info(f"Going to train on {device}")

train_set, valid_set, train_loader, valid_loader = mktrainval(args, logger)

logger.info(f"Loading model from {args.model}.npz")
model = models.KNOWN_MODELS[args.model](head_size=len(valid_set.classes), zero_head=True)
model.load_from(np.load(f"{args.model}.npz"))

logger.info("Moving model onto all GPUs")
model = torch.nn.DataParallel(model, device_ids=device_ids)

that just solve my problem

from big_transfer.