# Use vmap and grad to calculate gradients for one layer independently for each input in batch

**URL:** <https://discuss.pytorch.org/t/use-vmap-and-grad-to-calculate-gradients-for-one-layer-independently-for-each-input-in-batch/187556>\
**Category:** autograd\
**Created:** [September 1, 2023, 5:14pm UTC](https://discuss.pytorch.org/t/use-vmap-and-grad-to-calculate-gradients-for-one-layer-independently-for-each-input-in-batch/187556 "2023-09-01T17:14:40Z")\
**Posts on this page:** 1\
**Showing post:** 2

<div class="post-metadata">

**Author:** ![AlphaBetaGamma96](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/32.png) [@AlphaBetaGamma96](https://discuss.pytorch.org/u/AlphaBetaGamma96)\
**Post date:** [September 2, 2023, 9:51am UTC](https://discuss.pytorch.org/t/use-vmap-and-grad-to-calculate-gradients-for-one-layer-independently-for-each-input-in-batch/187556/2 "2023-09-02T09:51:29Z")

</div>

Hi @jdesmarais,

The reason why you don’t have a `.grad` attribute is that you’re mixing `torch.func` with `torch.autograd.grad`, which must be done with care.

If you want to vectorize your gradient calculation you’ll need to calculate it entirely within the `torch.func` namespace, so something like this,

```auto
def calc_loss(params, x): #and other args if need be
  return loss #insert loss func here

from torch.func import vmap
gradients = vmap(torch.func.grad(calc_loss, argnums=(0)), in_dims=(None,0))(params, x)

```

then you can concatenate them like you do above

---

_[View the full topic](https://discuss.pytorch.org/t/use-vmap-and-grad-to-calculate-gradients-for-one-layer-independently-for-each-input-in-batch/187556)._