fix: add missing s_dmask's gradient option

intelligent-machine-learning · Apr 17, 2024 · 52f1d1f · 52f1d1f
1 parent caaba0d
commit 52f1d1f
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -1045,7 +1045,8 @@ void InitXlaModuleBindings(py::module m) {
         if (return_softmax) {
           at::Tensor s_dmask =
               bridge::AtenFromXlaTensor(std::move(std::get<3>(result_tensors)));
-          result_tuple[3] = torch::autograd::make_variable(s_dmask, false);
+          result_tuple[3] =
+              torch::autograd::make_variable(s_dmask, query.requires_grad());
         }
         return result_tuple;
       },
@@ -1084,7 +1085,8 @@ void InitXlaModuleBindings(py::module m) {
         if (return_softmax) {
           at::Tensor s_dmask =
               bridge::AtenFromXlaTensor(std::move(std::get<3>(result_tensors)));
-          result_tuple[3] = torch::autograd::make_variable(s_dmask, false);
+          result_tuple[3] =
+              torch::autograd::make_variable(s_dmask, query.requires_grad());
         }
         return result_tuple;
       },