Fix the bug in training after FSDP eval

torch_xla/distributed/fsdp/xla_fully_sharded_data_parallel.py

@@ -1005,7 +1005,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
       # This can be used to debug FSDP parameter memory consumption.
       outputs = self._dummy_forward(*args, **kwargs)
 
-    if self.reshard_after_forward:
+    if self.reshard_after_forward or not torch.is_grad_enabled():
       output_opt_barrier_tensors = []
       if self.optimization_barrier_in_forward:
         # Ensure that the full parameters of this FSDP module are freed
@@ -1015,6 +1015,10 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
       self._free_full_params(
           dependency_tensors=output_opt_barrier_tensors,
           apply_opt_barrier=self.optimization_barrier_in_forward)
+      if not torch.is_grad_enabled():
+        for p in self.full_params:
+          if hasattr(p, '_param_infos'):
+            self.module.delete_unflatten_params_view(p._param_infos)
 
     # Register pre-backward hooks to all-gather the params for the backward
     # pass (if output's grad was needed). This won't register anything if