From 6d3570d0e2d9aa6af3d88a2f0640daed6fe6b433 Mon Sep 17 00:00:00 2001
From: bartosz roguski <bartosz.roguski@bluelabellabs.com>
Date: Wed, 12 Feb 2025 14:06:40 +0100
Subject: [PATCH] Add info on how to run tests to README.md

---
 README.md                 | 6 ++++++
 tests/test_consistency.py | 1 +
 2 files changed, 7 insertions(+)

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