This is indeed unrelated.
If you enable anomaly mode, you will see that the problem is that some of the params are saved for backward but modified inplace. The fix is to make sure they are not:
for i in range(epochs):
    model.train()
    train_loss = 0
    params = dict(model.named_parameters()) # !! Extract these for each inner loop
    for batch_idx, (x, y) in enumerate(dataset):
        params = {k: v.clone() for k,v in params.items()} # !! Make sure the modified params are not the ones you use
        logits = _stateless.functional_call(model, params, x)             # predict
        # ... rest of your code