Revert "Revert "Support torch.mean with dim=None (#3752)" (#3757)" (#3782)

* Revert "Revert "Support torch.mean with dim=None (#3752)" (#3757)"

This reverts commit ce1bd4e.

* Add torch pin

* Delete .torch_pin

Author: wonjoo-wj

torch_xla/csrc/aten_xla_type.cpp

@@ -1975,13 +1975,16 @@ at::Tensor XLANativeFunctions::mean(const at::Tensor& self,
       /*keep_reduced_dimensions=*/false, dtype));
 }
 
-at::Tensor XLANativeFunctions::mean(const at::Tensor& self, at::IntArrayRef dim,
-                                    bool keepdim,
+at::Tensor XLANativeFunctions::mean(const at::Tensor& self,
+                                    at::OptionalIntArrayRef dim, bool keepdim,
                                     c10::optional<at::ScalarType> dtype) {
   XLA_FN_COUNTER("xla::");
+  XLATensorPtr self_tensor = bridge::GetXlaTensor(self);
   return bridge::AtenFromXlaTensor(XLATensor::mean(
-      bridge::GetXlaTensor(self), torch::lazy::ToVector<int64_t>(dim),
-      /*keep_reduced_dimensions=*/keepdim, dtype));
+      self_tensor,
+      dim ? torch::lazy::ToVector<int64_t>(*dim)
+          : torch::lazy::Iota<int64_t>(self_tensor->shape().get().rank()),
+      keepdim, dtype));
 }
 
 at::Tensor XLANativeFunctions::min(const at::Tensor& self) {