Creating a tensor on the GPU freezes the process

username: Sa96
cooked: <p>I am running <a href="https://github.com/haqishen/SIIM-ISIC-Melanoma-Classification-1st-Place-Solution" rel="noopener nofollow ugc">this repo</a> on my system ( MIG device , 80GB of memory) to train a classifier. I want to pass a list of weights to the loss function in train.py as below:</p>
<pre><code class="lang-auto">if __name__ == '__main__':

    args = parse_args()
    os.makedirs(args.model_dir, exist_ok=True)
    os.makedirs(args.log_dir, exist_ok=True)
    os.environ['CUDA_VISIBLE_DEVICES'] = args.CUDA_VISIBLE_DEVICES

    if args.enet_type == 'resnest101':
        ModelClass = Resnest_Melanoma
    elif args.enet_type == 'seresnext101':
        ModelClass = Seresnext_Melanoma
    elif 'efficientnet' in args.enet_type:
        ModelClass = Effnet_Melanoma
    else:
        raise NotImplementedError()

    DP = len(os.environ['CUDA_VISIBLE_DEVICES']) &gt; 1

    set_seed()
    
    device = torch.device('cuda')

    c_weights = torch.tensor([0.954, 0.8274, 0.852, 0.987, 0.967, 0.986, 0.735, 0.687]).float().to(device)
    #loss function
    criterion = nn.CrossEntropyLoss()
    main()
</code></pre>
<p>When I create the tensor on CPU it works just fine (I need to use a GPU because of my image). But when I send the tensor <code>c_weights</code> to the GPU the train process gets stuck without returning any errors. I tried different methods for creating the tensor and different dtypes it did not work. This is what it shows in the output (the process gets stuck on 2 and does not progress):</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/1/e1b2ab75dd6483ac10ae9565c0813385088df2da.png" data-download-href="https://discuss.pytorch.org/uploads/default/e1b2ab75dd6483ac10ae9565c0813385088df2da" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e1b2ab75dd6483ac10ae9565c0813385088df2da_2_690x30.png" alt="image" data-base62-sha1="wcCecvR8rev5FXKktussOmAGUdA" width="690" height="30" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e1b2ab75dd6483ac10ae9565c0813385088df2da_2_690x30.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e1b2ab75dd6483ac10ae9565c0813385088df2da_2_1035x45.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/1/e1b2ab75dd6483ac10ae9565c0813385088df2da_2_1380x60.png 2x" data-dominant-color="161616"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">2544×114 8.95 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I tried smaller batch sizes as well. No improvement.<br>
Here is my GPU information:</p>
<p><em>GPU Driver Version: 510.85.02</em><br>
<em>CUDA Version: 11.6</em><br>
<em>memory: 81069MiB</em></p>
<p>How can I solve this issue?</p>
<p><strong>Edit:</strong><br>
My args are:</p>
<pre><code class="lang-auto">--kernel-type = 8c_b3_768_512_18ep
--data-folder = 512
--image-size = 512
--enet-type = efficientnet_b3
--batch-size = 32
--num-workers = 32
--out-dim = 8
--CUDA_VISIBLE_DEVICES = MIG-cdc1351f-1b7a-554c-a273-f7643f99523f
--fold = 0

</code></pre>

username: tslavik
cooked: <p>I was able to get the device and c_weights lines to run as is.  I would first verify args.CUDA_VISIBLE_DEVICES is correct then try adding a couple of print lines for device and visible devices, just before c_weights, to see if that might be the issue.</p>

username: Sa96
cooked: <p>I edited my post and added the args I’m using. I printed these info:</p>
<pre><code class="lang-auto">cuda visible device : MIG-cdc1351f-1b7a-554c-a273-f7643f99523f
device :  cuda
c_weights: tensor([0.9540, 0.8274, 0.8520, 0.9870, 0.9670, 0.9860, 0.7350, 0.6870],
       device='cuda:0')
</code></pre>