[autograd] refactor some functions (#37061)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/37061
This PR refactors:
1. `set_device` to make it out of Engine
2. put `graph_task_completed` into GraphTask
3. put `mark_graph_task_completed` into GraphTask
This also make the distributed engine easy to call those functions.
Test Plan: Imported from OSS
Differential Revision: D21188688
Pulled By: wanchaol
fbshipit-source-id: f56106e6ed7d966cfa4d962781c7865cc3c5321d
