-
Notifications
You must be signed in to change notification settings - Fork 3.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
move device-specific teardown logic from training loop to accelerator #5973
Conversation
Codecov Report
@@ Coverage Diff @@
## master #5973 +/- ##
======================================
- Coverage 90% 90% -0%
======================================
Files 170 170
Lines 11789 11786 -3
======================================
- Hits 10669 10637 -32
- Misses 1120 1149 +29 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM !
@@ -27,6 +27,7 @@ def on_train_start(self): | |||
|
|||
def on_train_end(self): | |||
# clean up memory | |||
self.model.cpu() |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we do this for TPU too ?
self.trainer.accelerator_backend.on_train_end() | ||
|
||
# clear mem | ||
if self.trainer._device_type == DeviceType.GPU: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nice cleaning !
What does this PR do?
Follow up to #5743
on_train_end device-specific teardown should be handled by accelerator.
Makes training loop device agnostic
Before submitting
PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:
Did you have fun?
Make sure you had fun coding 🙃