In most codes the order I see is
training loop:
    # forward pass and calculate loss
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

If I change it to:
training loop:
    # forward pass and calculate loss
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

is it still ok?