How could I train on multi-gpu and infer with single gpu

5

I have access to my gpus, the program works when I run python infer.py, but it will not work if I run CUDA_VISIBLE_DEVICES python infer.py.
The root of this problem seems to be that I train my model with two gpus (nn.DataParallel), but I run test on a single gpu.