Update functorch supported autograd.Function to allow mark_dirty (#91222)
Fixes https://github.com/pytorch/pytorch/issues/90225
Uses what was originally in https://github.com/pytorch/pytorch/pull/89860/commit/32a57bcdb6328c6f46a4908d96d652ab539fdf83
Pull Request resolved: https://github.com/pytorch/pytorch/pull/91222
Approved by: https://github.com/zou3519
