(params, buffers, sample, target)
| 135 | # use ``torch.unsqueeze`` to add a batch dimension. |
| 136 | |
| 137 | def compute_loss(params, buffers, sample, target): |
| 138 | batch = sample.unsqueeze(0) |
| 139 | targets = target.unsqueeze(0) |
| 140 | |
| 141 | predictions = functional_call(model, (params, buffers), (batch,)) |
| 142 | loss = loss_fn(predictions, targets) |
| 143 | return loss |
| 144 | |
| 145 | ###################################################################### |
| 146 | # Now, let’s use the ``grad`` transform to create a new function that computes |