diff options
Diffstat (limited to 'torch')
| -rw-r--r-- | torch/csrc/autograd/aten_variable_hooks.cpp | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/torch/csrc/autograd/aten_variable_hooks.cpp b/torch/csrc/autograd/aten_variable_hooks.cpp index b1ef0c948c..7a2c3974c2 100644 --- a/torch/csrc/autograd/aten_variable_hooks.cpp +++ b/torch/csrc/autograd/aten_variable_hooks.cpp @@ -16,7 +16,8 @@ REGISTER_VARIABLE_HOOKS(VariableHooks) // Pre-condition: backend/scalar_type is a valid type in the type_registry void VariableHooks::registerVariableTypeFor(at::Context* context, at::Backend backend, at::ScalarType scalar_type) const { - register_variable_type_for(context, backend, scalar_type); + auto* baseType = context->getTypeRaw(backend, scalar_type); + register_variable_type_for(baseType); } }} // torch::autograd |