diff --git a/README.md b/README.md
index e2f4206..abda088 100644
--- a/README.md
+++ b/README.md
@@ -221,6 +221,12 @@ docker run -it --runtime=habana clip_hpu:latest
 
 You do not need to change the code to leverage Intel® Gaudi® HPU. The `get_device_initial()` function will automatically detect the correct device and return the appropriate device name. So no changes are required.
 
+### Run the Tests
+
+```bash
+pytest
+```
+This will run the tests and verify that the model is working correctly.
 
 ## See Also
 
diff --git a/tests/test_consistency.py b/tests/test_consistency.py
index 3717250..9b72208 100644
--- a/tests/test_consistency.py
+++ b/tests/test_consistency.py
@@ -2,6 +2,7 @@ import numpy as np
 import pytest
 import torch
 from PIL import Image
+import habana_frameworks.torch
 
 import clip