Skip to content

Fix torch.jit.ScriptModule.zero_grad. #1478

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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

hiyuh
Copy link
Contributor

@hiyuh hiyuh commented May 30, 2025

TorchSharp 0.105.0 doesn't have torch.jit.ScriptModule.zero_grad and falls back into torch.nn.Module.zero_grad incorrectly, then terminates silently.
Most probably, because JITModule is not compatible to NNModule in LibTorchSharp.

And as reported in pytorch/pytorch#27144, libtorch also doesn't have torch::jit::Module::zero_grad.
As a workaround, manually loop over the parameters and zero them out like optimizer does.

Note;

  • intentionally, omit RELEASENOTES.md update ATM.
    • due to avoid multiple conflict&rebase annoyance while MR review.
    • i'll update later, before merging this MR, if upstream prefers.
  • i'm not sure whether foreach loop of ScriptModule.zero_grad in src/TorchSharp/JIT/ScriptModule.cs is actually needed.
    • this is just mimicking what Module.zero_grad in src/TorchSharp/NN/Module.cs does.

TorchSharp 0.105.0 doesn't have torch.jit.ScriptModule.zero_grad and
falls back into torch.nn.Module.zero_grad incorrectly, then terminates
silently.
Most probably, because JITModule is not compatible to NNModule in
LibTorchSharp.

And as reported in pytorch/pytorch#27144,
libtorch also doesn't have torch::jit::Module::zero_grad.
As a workaround, manually loop over the parameters and zero them out
like optimizer does.
@hiyuh hiyuh force-pushed the fix-ScriptModule-zero_grad branch from 35aa418 to dc6fc3c Compare May 30, 2025 06:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant