Why do we need to do loss.cuda() when we we have already done model.cuda()?

username: samra-irshad
cooked: <p>I am following a source code corresponding to a paper where authors have trained the model to do a segmentation task. In the code, authors wrote:</p>
<pre><code class="lang-auto"> if torch.cuda.is_available():
        net.cuda()
        softMax.cuda()
        CE_loss.cuda()
        Dice_loss.cuda()
</code></pre>
<p>I don’t understand why the authors felt the need to send the loss functions to CUDA when the model has already been sent to the device? I have been using tensorflow in my past and I am new to pytorch, so I have trouble in understanding the device initialization for models and loss functions.</p>

username: Nikronic
cooked: <p>Hi,</p>
<p>I do not know the implementation of those loss functions exactly (if they are modified), but if a criterion does not have any parameters, then sending to cuda makes no difference as there is no operation to change parameters in cuda.</p>
<aside class="quote" data-post="2" data-topic="85716">
  <div class="title">
    <div class="quote-controls"></div>
    <img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/nikronic/48/71873_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/is-that-ok-not-using-criterion-cuda-before-criterion/85716/2">Is that ok not using 'criterion.cuda()' before criterion?</a> 
  </div>
  <blockquote>
    Hi, 
If I want to summarize, you only need to use .cuda on criterion if the criterion function has parameters. 
Here is the posts related to this situation: <a href="https://discuss.pytorch.org/t/move-the-loss-function-to-gpu/20060" class="inline-onebox">Move the loss function to GPU</a> and <a href="https://discuss.pytorch.org/t/why-criterion-cuda-is-not-needed-but-model-cuda-is/17410/2" class="inline-onebox">Why `criterion.cuda()` is not needed but `model.cuda()` is?</a> 
Bests
  </blockquote>
</aside>

<p>Bests</p>

username: samra-irshad
cooked: <p>What do you mean by ‘if the criterion does not have any parameters’? So far, I have only seen non-parameterized loss functions, i.e., loss functions which get some input tensors, apply loss function on those tensors and then give outputs.</p>
<p>Do you mean if the loss function being applied to tensors? Because I can imagine the loss function is being applied to CUDA tensors.</p>

username: Nikronic
cooked: <p>Yes, exactly, parametrized loss functions are rare and that is why you cannot find <code>loss.cuda</code> anywhere.</p>
<aside class="quote no-group" data-username="samra-irshad" data-post="3" data-topic="91023">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/samra-irshad/48/27993_2.png" class="avatar"> samra-irshad:</div>
<blockquote>
<p>loss functions which get some input tensors, apply loss function on those tensors and then give outputs.</p>
</blockquote>
</aside>
<p>Here, no need to use <code>loss.cuda</code> as its inputs all are cuda tensors.</p>

username: samra-irshad
cooked: <p>Right. So as described in my original post, softmax function has also been moved to CUDA. I briefly reviewed the source code for <code>nn.Softmax()</code> and saw these lines:</p>
<pre><code class="lang-auto">class Softmax(Module):
  
    __constants__ = ['dim']
    dim: Optional[int]

    def __init__(self, dim: Optional[int] = None) -&gt; None:
        super(Softmax, self).__init__()
        self.dim = dim

    def __setstate__(self, state):
        self.__dict__.update(state)
        if not hasattr(self, 'dim'):
            self.dim = None

    def forward(self, input: Tensor) -&gt; Tensor:
        return F.softmax(input, self.dim, _stacklevel=5)

    def extra_repr(self) -&gt; str:
        return 'dim={dim}'.format(dim=self.dim)
</code></pre>
<p>I am not sure why<code> def__setstate__</code> is used and what state is it setting? How does it receive the state? Also, I noticed the module/function that inherits nn.Module class usually gets to move to CUDA. Any thoughts?</p>