diff --git a/src/pystencils_autodiff/backends/_torch_native.py b/src/pystencils_autodiff/backends/_torch_native.py index fa5c2997817b91383fe5c9b768235ab110fe8e03..bba20bc7417c41765f16c10ea6193cfdf7245b25 100644 --- a/src/pystencils_autodiff/backends/_torch_native.py +++ b/src/pystencils_autodiff/backends/_torch_native.py @@ -73,6 +73,8 @@ def create_autograd_function(autodiff_obj, use_cuda): output_tensors = OrderedDict({f.name: field_to_tensor_dict.get(f, kwargs[f.name]) for f in autodiff_obj.forward_output_fields}) + + # TODO: only if any of input tensors requires grad for o in output_tensors.values(): if isinstance(o, torch.Tensor): o.requires_grad = True