Skip to content
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 criterion to compute device automatically #455

Merged
merged 3 commits into from Apr 15, 2019

Conversation

Projects
None yet
2 participants
@ottonemo
Copy link
Collaborator

commented Apr 11, 2019

Previously the user had to make sure that criterion parameters
such as class weights are on the correct computing device.
Now the criterion is moved to the compute device using .to
which moves the parameters as well.

Move criterion to compute device automatically
Previously the user had to make sure that criterion parameters
such as class weights are on the correct computing device.
Now the criterion is moved to the compute device using `.to`
which moves the parameters as well.

@ottonemo ottonemo requested a review from BenjaminBossan Apr 11, 2019

@ottonemo

This comment has been minimized.

Copy link
Collaborator Author

commented Apr 11, 2019

One issue with this is that we would then force criterion to be a subclass of torch.nn.Module since whatever is passed needs to have the .to method. We could live with that or add something like

        self.criterion_ = self.criterion(**criterion_params)
        if isinstance(self.criterion_, torch.nn.Module):
            self.criterion_ = self.criterion_.to(self.device)

What do you think?

@BenjaminBossan

This comment has been minimized.

Copy link
Collaborator

commented Apr 11, 2019

Yes, I believe we should allow any kind of nn.Module. Furthermore, we might add a note to the docstring of criterion that tells the user that if they wish to have proper device handling on the criterion, it should be an instance of nn.Module.

@ottonemo

This comment has been minimized.

Copy link
Collaborator Author

commented Apr 12, 2019

@BenjaminBossan Fixed, please review again.

@BenjaminBossan BenjaminBossan merged commit 81ae225 into master Apr 15, 2019

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.