You can do per example gradients using torch.autograd.grad.  The key trick is to manually broadcast the weights yourself, and then differentiate with respect to the broadcasted weight. Here’s a very simple example on a model that is just a Linear:
import torch

torch.manual_seed(0)

B = 4
N = 2
M = 3

input = torch.randn(B, 1, N)
weights = torch.randn(N, M, requires_grad=True)

broadcasted_weights = weights.expand(B, N, M)
output = torch.bmm(input, broadcasted_weights)
loss = output.sum()

print(torch.autograd.grad(loss, broadcasted_weights))

I don’t know if you can conveniently do this with torch.nn modules, might be worth thinking about how to expose this in the API.