[Examples] Update examples to record total loss via `total_loss += loss.detach()`
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Most of the current examples have a training loops like one of the following:
Using `loss.item()`:
```python3
for x, y in dataloader:
y_pred = model(x)
loss = loss_fn(y, y_pred)
opt.zero_grad()
loss.backward()
opt.step()
total_loss += loss.item()
```
Calling `loss.item()` requires synchronizing the CPU and GPU, and prevents overlapping forward/backward/optimizer with the dataloader. See the Nsight Systems profile below:

Using `loss` directly will avoid this:
```python3
for x, y in dataloader:
y_pred = model(x)
loss = loss_fn(y, y_pred)
opt.zero_grad()
loss.backward()
opt.step()
total_loss += loss
```
However, while in many cases this will not negatively impact the computation, it can end up connecting `total_loss` to the computation graph.
What we should be doing is calling `loss.detach()`:
```python3
for x, y in dataloader:
y_pred = model(x)
loss = loss_fn(y, y_pred)
opt.zero_grad()
loss.backward()
opt.step()
total_loss += loss.detach()
```
This accumulates the loss as a 1-element tensor on the training device, not attached to the computation graph, and allows the CPU to move on to the dataloader when the GPU is still working on forward/backward/optimizer.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.