Repository: NVIDIA-Merlin/NVTabular

Issue:
[BUG] TF CPU dataloader uses GPU memory if GPU is available 
**Describe the bug**
I'm getting "cudaErrorMemoryAllocation out of memory", when trying to use KerasSequenceLoader with CPU device, which means that the dataloader is trying to access the GPU memory even when setting device to be CPU.

**Steps/Code to reproduce bug**
Run any model script in the background. for example `python model.py`
```python
import tensorflow as tf

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])

input() # Prevent script from finishing
```

When running the `model.py` in the background (to keep GPU memory allocated by tensorflow), run `python dataloader.py`

```python
from nvtabular.loader.tensorflow import KerasSequenceLoader

LABEL=...
PARQUET_PATH=...
NUMERIC_COLUMNS=...
CATEGORICAL_COLUMNS=...

train_dataset_tf = KerasSequenceLoader(
    PARQUET_PATH,
    batch_size=4096,
    label_names=LABEL,
    cat_names=CATEGORICAL_COLUMNS,
    cont_names=NUMERIC_COLUMNS,
    engine="parquet",
    shuffle=False,
    buffer_size=0,
    parts_per_chunk=1,
    device='cpu',
    reader_kwargs={"cpu": True}
)

for x in train_dataset_tf:
    print(x)
```

which will yield
```
MemoryError: std::bad_alloc: CUDA error at: ../include/rmm/mr/device/cuda_memory_resource.hpp:70: cudaErrorMemoryAllocation out of memory
```

**Expected behavior**
No GPU memory is required to run KerasSequenceLoader with device='cpu', even if the gpu device is available.

**Environment details (please complete the following information):**
 - Environment location: `nvcr.io/nvidia/merlin/merlin-tensorflow-training:22.02`
 



Discussion:

Resolving patch:
diff --git a/nvtabular/loader/tensorflow.py b/nvtabular/loader/tensorflow.py
--- a/nvtabular/loader/tensorflow.py
+++ b/nvtabular/loader/tensorflow.py
@@ -32,7 +32,7 @@
 
 # noqa
 try:
-    from nvtabular.io import Dataset
+    from merlin.io import Dataset
 
     nvt_dataset_class = Dataset
 except ImportError:
@@ -49,7 +49,7 @@
 }
 
 
-def _validate_dataset(paths_or_dataset, batch_size, buffer_size, engine, reader_kwargs):
+def _validate_dataset(paths_or_dataset, batch_size, buffer_size, engine, device, reader_kwargs):
     # TODO: put this in parent class and allow
     # torch dataset to leverage as well?
 
@@ -76,8 +76,11 @@ def _validate_dataset(paths_or_dataset, batch_size, buffer_size, engine, reader_
     if not engine:
         # default engine is parquet
         engine = "parquet"
+
+    cpu = device and "cpu" in device
+
     if nvt_dataset_class:
-        return nvt_dataset_class(files, engine=engine)
+        return nvt_dataset_class(files, engine=engine, cpu=cpu)
     else:
         LOG.warning(
             "NVTabular Dataset class not detected, reverting to Dask Dataframe."
@@ -252,16 +255,16 @@ def __init__(
         sparse_as_dense=False,
         schema=None,
     ):
+        device = device or 0
+        device = "cpu" if not HAS_GPU else device
         dataset = _validate_dataset(
-            paths_or_dataset, batch_size, buffer_size, engine, reader_kwargs
+            paths_or_dataset, batch_size, buffer_size, engine, device, reader_kwargs
         )
         schema = _get_schema(dataset) if not schema else schema
         cat_names, cont_names = _validate_schema(
             feature_columns, cat_names, cont_names, schema=schema
         )
 
-        device = device or 0
-        device = "cpu" if not HAS_GPU else device
         DataLoader.__init__(
             self,
             dataset,


Test patch: