multi gpu error

CUDA_VISIBLE_DEVICES=6,7 python train.py --PCB --batchsize 60 --name PCB-64 --train_all

I use multi gpu, so add some code:
if torch.cuda.device_count() > 1 and use_gpu:
  model_wraped = nn.DataParallel(model).cuda()
  model = model_wraped

but error in forward: RuntimeError: cuda runtime error (77) : an illegal memory access was encountered at /opt/conda/conda-bld/pytorch_1513368888240/work/torch/lib/THC/THCTensorCopy.cu:204

[layumi]

This is my code. You may refer it and modify your code.

model = torch.nn.DataParallel(model, device_ids=gpu_ids).cuda()

ignored_params = list(map(id, model.module.model.fc.parameters() )) + list(map(id, model.module.classifier.parameters() ))
base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())

# Observe that all parameters are being optimized
optimizer_ft = optim.SGD([
             {'params': base_params, 'lr': 0.01},
             {'params': model.module.model.fc.parameters(), 'lr': 0.1},
             {'params': model.module.classifier.parameters(), 'lr': 0.1}
         ], momentum=0.9, weight_decay=5e-4, nesterov=True)

[xujian0]

Hi @xujian0 When you test the model, you also need to add model.module in the test.py. Especially, this line

model = torch.nn.DataParallel(model, device_ids=gpu_ids).cuda()

Thanks for your reply, and it solved my problem!