If your full model has model.part1 and model.part2, then you can do:
grad_part1 = torch.autograd.grad(loss, model.part1.parameters(), retain_graph=True)
grad_part2 = torch.autograd.grad(loss, model.part2.parameters())

Does that match what you want?