fix: set use_reentrant to True to fix Mixtral-7b bug

[geoffreyangus]

Got the following error when training Mixstral-7b in a multi-GPU setting:

tensor at position 96:
saved metadata: {'shape': torch.Size([142]), 'dtype': torch.int64, 'device': device(type='cuda', index=1)}
recomputed metadata: {'shape': torch.Size([143]), 'dtype': torch.int64, 'device': device(type='cuda', index=1)}
tensor at position 97:
saved metadata: {'shape': torch.Size([142]), 'dtype': torch.int64, 'device': device(type='cuda', index=1)}
recomputed metadata: {'shape': torch.Size([143]), 'dtype': torch.int64, 'device': device(type='cuda', index=1)}
...

This can be resolved by setting reentrant to True ensures

Reentrant checkpoint always recomputes function in its entirety during the backward pass.
Source: https://pytorch.org/docs/stable/checkpoint.html

So shape mismatches should be prevented.

[arnavgarg1]

Can we also bump the minimum torch version as a follow up PR?