Fix potential memory leak (#1167)

Lightning-Universe · Feb 14, 2022 · 381aa37 · 381aa37
1 parent 3ec0405
commit 381aa37
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flash/core/data/io/input.py b/flash/core/data/io/input.py
@@ -14,7 +14,7 @@
 import functools
 import os
 import sys
-from copy import copy, deepcopy
+from copy import deepcopy
 from dataclasses import dataclass
 from functools import partial
 from typing import Any, Callable, cast, Dict, Iterable, List, MutableMapping, Optional, Sequence, Tuple, Type, Union
@@ -233,7 +233,9 @@ def _call_load_sample(self, sample: Any) -> Any:
                 InputBase,
             ),
         )
-        return load_sample(copy(sample))
+
+        # Deepcopy the sample to avoid leaks with complex data structures
+        return load_sample(deepcopy(sample))
 
     @staticmethod
     def load_data(*args: Any, **kwargs: Any) -> Union[Sequence, Iterable]: