device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if opt.resume:
    if os.path.isfile(opt.resume):
        print("Loading from checkpoint {}".format(opt.resume))
        model = torch.load(opt.resume)
        model.load_state_dict(model.state_dict())
        opt.start_training_step, opt.start_epoch = which_trainingstep_epoch(opt.resume)
else:
    model = Net()
    mkdir_steptraing(opt.model_name)

model = model.to(device)
model = nn.DataParallel(model, device_ids=range(opt.n_GPUs))