🐛 Describe the bug

Summary

torch.jit.ScriptModule.register_full_backward_hook() raises RuntimeError: Could not cast attribute '_is_full_backward_hook' to type NoneType: Cannot cast True to None when called on a ScriptModule subclass, while the same API works correctly on a plain nn.Module.

Reproduction

import torch
import torch.nn as nn


class MyModel(torch.jit.ScriptModule):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 2)

    def forward(self, x):
        return self.linear(x)


model = MyModel()
model.register_full_backward_hook(lambda module, grad_input, grad_output: grad_input)
# RuntimeError: Could not cast attribute '_is_full_backward_hook' to type NoneType: Cannot cast True to None

Expected Behavior

Should work the same as on nn.Module — register the hook and return a RemovableHandle.

Root Cause Analysis

ScriptModule uses a two-layer structure:

• Outer layer: the user's class (e.g. MyModel), which inherits from ScriptModule -> nn.Module, with Python-side methods available
• Inner layer: _actual_script_module, a RecursiveScriptModule instance created during init_then_script, which wraps the compiled C++ TorchScript module

The call chain is:

1. register_full_backward_hook() (defined in nn.Module) executes self._is_full_backward_hook = True
2. This triggers ScriptModule.__setattr__, which unconditionally delegates to setattr(self._actual_script_module, attr, value) since _actual_script_module already exists in __dict__
3. This triggers RecursiveScriptModule.__setattr__, which finds _c.hasattr('_is_full_backward_hook') is True (the attribute was registered as None during nn.Module.__init__ before compilation)
4. self._c.setattr('_is_full_backward_hook', True) is called on the C++ side, which expects the attribute type to remain NoneType, causing the cast failure

The core issue is that ScriptModule.__setattr__ delegates all attribute assignments to the inner RecursiveScriptModule without distinguishing which attributes should stay on the Python side. _is_full_backward_hook is a Python-side runtime flag (used to enforce mutual exclusion with the deprecated register_backward_hook), it should be set on the Python __dict__ rather than forwarded to the C++ module.

In contrast, register_full_backward_pre_hook works correctly because it does not set _is_full_backward_hook internally.

Suggested Fix

ScriptModule.__setattr__ should keep Python-side runtime attributes (like _is_full_backward_hook, _backward_hooks, _backward_pre_hooks, etc.) in the Python __dict__ instead of delegating them to the C++ RecursiveScriptModule. One approach is to check whether the attribute is a Python-only hook-related attribute before delegating.

Versions

• PyTorch version: 2.11
• Python version: 3.12

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel