How to run knn2img.py on multiple GPUs?

[THEWEAKEST]

I use 4 TITAN Xp each of them has 12GB memory . And I modify the knn2img.py

    config = OmegaConf.load(f"{opt.config}")
    model = load_model_from_config(config, f"{opt.ckpt}")

    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = torch.nn.DataParallel(model, device_ids = [0,1,2,3]) # add
    model = model.to(device)

instead of

    config = OmegaConf.load(f"{opt.config}")
    model = load_model_from_config(config, f"{opt.ckpt}")

    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
    model = model.to(device)

Finally I use CUDA_VISIBLE_DEVICES = 0,1,2,3 But the model still runs on GPU0 without using other GPUs and then cause CUDA out of Memory . Is there any way to run knn2img.py on multiple GPUs ? Thanks very much .

[GUOENHUA]

Have you ever solved this issue?