CrossEntropyLoss doesn't move the weight to same device

glenguo06: <p>It appears that for CrossEntropyLoss, <code>.to()</code> doesn’t copy the <code>weight</code> to the same device. Would be great to know is this by-design or a bug that is fixed on newer releases?</p>
<p>Minimum example from the usage doc:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; loss = nn.CrossEntropyLoss(weight=torch.FloatTensor([1, 2, 3])).cuda()
&gt;&gt;&gt; input = torch.randn(3, 5, requires_grad=True).cuda()
&gt;&gt;&gt; target = torch.empty(3, dtype=torch.long).random_(5).cuda()
&gt;&gt;&gt; output = loss(input, target)
</code></pre>
<p>It will throw the following exception:</p>
<pre><code class="lang-auto">RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0!
</code></pre>
<p>Looking at the <code>_WeightedLoss</code> module, <code>weight</code> parameter is registered into the buffer. I assume it should be copied to same device when <code>.to()</code> is called from the loss module?<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/0/10ccb444e94588fc96239b1083cce2b48800bc90.png" data-download-href="https://discuss.pytorch.org/uploads/default/10ccb444e94588fc96239b1083cce2b48800bc90" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/0/10ccb444e94588fc96239b1083cce2b48800bc90.png" alt="image" data-base62-sha1="2oCdp58WfNeyQZid5fKXGGaCtNe" width="690" height="88" data-dominant-color="232424"><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">1266×162 10.7 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>

KFrank: <p>Hi Glen!</p>
<aside class="quote no-group" data-username="glenguo06" data-post="1" data-topic="216873">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/glenguo06/48/25474_2.png" class="avatar"> glenguo06:</div>
<blockquote>
<p>It appears that for CrossEntropyLoss, <code>.to()</code> doesn’t copy the <code>weight</code> to the same device. Would be great to know is this by-design or a bug that is fixed on newer releases?</p>
</blockquote>
</aside>
<p>It works for me on version 2.6.0:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.6.0+cu126
&gt;&gt;&gt; 
&gt;&gt;&gt; ce = torch.nn.CrossEntropyLoss (weight = torch.ones (3))
&gt;&gt;&gt; 
&gt;&gt;&gt; ce.weight
tensor([1., 1., 1.])
&gt;&gt;&gt; list (ce.buffers())
[tensor([1., 1., 1.])]
&gt;&gt;&gt; _ = ce.cuda()
&gt;&gt;&gt; ce.weight
tensor([1., 1., 1.], device='cuda:0')
&gt;&gt;&gt; 
&gt;&gt;&gt; input = torch.randn (3, device = 'cuda')
&gt;&gt;&gt; target = torch.tensor (1, device = 'cuda')
&gt;&gt;&gt; ce (input, target)
tensor(0.3248, device='cuda:0')
</code></pre>
<p><code>weight</code> is a <code>buffer</code> of the <code>Module</code> <code>CrossEntropyLoss</code>, so it is expected to work.</p>
<p>Best.</p>
<p>K. Frank</p>