diff --git a/merlin/core/compat/__init__.py b/merlin/core/compat/__init__.py index 8ff3eede1..4d81d60b9 100644 --- a/merlin/core/compat/__init__.py +++ b/merlin/core/compat/__init__.py @@ -101,7 +101,6 @@ def device_mem_size(kind="total", cpu=False): return pynvml_mem_size(kind=kind) - try: import numpy except ImportError: diff --git a/merlin/io/writer.py b/merlin/io/writer.py index 2acd75884..c4a5aa788 100644 --- a/merlin/io/writer.py +++ b/merlin/io/writer.py @@ -196,7 +196,11 @@ def _add_data_slice(self, df): if self.shuffle: df = shuffle_df(df) int_slice_size = df.shape[0] // self.num_out_files - slice_size = int_slice_size if int_slice_size > 0 and df.shape[0] % int_slice_size == 0 else int_slice_size + 1 + slice_size = ( + int_slice_size + if int_slice_size > 0 and df.shape[0] % int_slice_size == 0 + else int_slice_size + 1 + ) for x in range(self.num_out_files): start = x * slice_size end = start + slice_size