diff --git a/README.md b/README.md
index 72b15004f9..38ea969962 100644
--- a/README.md
+++ b/README.md
@@ -86,10 +86,12 @@ td.save(feature_join)
 
 Feed the training dataset to a TensorFlow model:
 ```python
-train_input_feeder = training_dataset.feed(target_name="label",
-                                           split="train",
-                                           is_training=True)
-train_input = train_input_feeder.tf_record_dataset()
+tf_data_object = training_dataset.tf_data(target_name="label",
+                                          split="train",
+                                          is_training=True)
+train_input = tf_data_object.tf_record_dataset(batch_size=32,
+                                               num_epochs=5,
+                                               process=True)
 ```
 
 You can find more examples on how to use the library in our [hops-examples](https://github.com/logicalclocks/hops-examples) repository.
diff --git a/python/hsfs/core/tfdata_engine.py b/python/hsfs/core/tfdata_engine.py
index 7990c6ddf7..fad28ae0c0 100644
--- a/python/hsfs/core/tfdata_engine.py
+++ b/python/hsfs/core/tfdata_engine.py
@@ -16,17 +16,18 @@
 
 import re
 import itertools
+import mock
 from typing import Optional
 
 try:
     import tensorflow as tf
 except ModuleNotFoundError:
-    pass
+    tf = mock.Mock()
 
 try:
     from pydoop import hdfs
 except ModuleNotFoundError:
-    pass
+    hdfs = mock.Mock()
 
 try:
     import boto3
diff --git a/python/setup.py b/python/setup.py
index b8a72f9265..208543d8e7 100644
--- a/python/setup.py
+++ b/python/setup.py
@@ -26,6 +26,7 @@ def read(fname):
         "PyMySQL",
         "pyjks",
         "sqlalchemy",
+        "mock",
     ],
     extras_require={
         "dev": [