rm rm (NVIDIA#11116)

Signed-off-by: Alexandros Koumparoulis <[email protected]>
Signed-off-by: Hainan Xu <[email protected]>

nemo/lightning/io/pl.py

@@ -142,21 +142,13 @@ def save_checkpoint(self, checkpoint: Dict[str, Any], path: _PATH, storage_optio
         validate_sharding_integrity = not (self.validated_consistency and self.assume_constant_structure)
         self.validated_consistency = True
 
-        try:
-            return dist_checkpointing.save(
-                sharded_state_dict=checkpoint,
-                checkpoint_dir=checkpoint_dir,
-                sharded_strategy=self.save_sharded_strategy,
-                validate_access_integrity=validate_sharding_integrity,
-                async_sharded_save=self.async_save,
-            )
-        except:
-            logging.error(f"Failed to save checkpoint to {checkpoint_dir}")
-            # Do cleanup.
-            import shutil
-
-            shutil.rmtree(checkpoint_dir)
-            raise
+        return dist_checkpointing.save(
+            sharded_state_dict=checkpoint,
+            checkpoint_dir=checkpoint_dir,
+            sharded_strategy=self.save_sharded_strategy,
+            validate_access_integrity=validate_sharding_integrity,
+            async_sharded_save=self.async_save,
+        )
 
     @override
     def load_checkpoint(