Fix: Preserve model device in model_summary method#574
Fix: Preserve model device in model_summary method#574rwood-97 merged 4 commits intomaps-as-data:mainfrom
Conversation
|
Hi, the tests are failing but I think this is a problem on my side. I will fix this in the MapReader repo and then you will need to sync the changes from MapReader into your fork and then rebase your branch using |
|
Hi @GiuseTripodi , I have (hopefully) now fixed the tests in the main MapReader repo, can you do the steps above and then let me know when you've done it and I'll review/approve the PR. The changelog test checks whether you've updated the changelog. You'll need to add a |
Added device to the summary function to cleary define where the model and data needs to be stored.
5c6dbfc to
03847df
Compare
|
Hi @rwood-97, I have done the steps above. Please let me know if there is anything else I should do. |
| --- | ||
|
|
||
| ## Pre-release | ||
| _Add new changes here_ |
There was a problem hiding this comment.
Can you re-add the _Add new changes here_ line so future contributors can still see where to add in the changelog
There was a problem hiding this comment.
Sure, done it.
rwood-97
left a comment
There was a problem hiding this comment.
Great, thanks for this.
I've added one comment re keeping the _Add new changes here_ line in the changelog. Otherwise looks good to me.
I've also noted that one of the tests is still failing but I think thats a separate issue (i.e. not for you to fix in this PR). I've created an issue on the repo to track it -> #576
Summary
The 'model_summary()' method was resetting the model device from MPS or CUDA back to CPU after calling 'torchinfo.summary()'. By adding explicitly 'device' parameter to torchinfo.summar() the model stays on its original device throughout the summary operation.
Fixes #542
Describe your changes
Modified 'model_summary()' method by adding the 'device' parameter and setting it to the current device.
Checklist before assigning a reviewer (update as needed)
Reviewer checklist
Please add anything you want reviewers to specifically focus/comment on.