From a2e083b3150edb0846a9eb5f1dd7129db61c6e05 Mon Sep 17 00:00:00 2001
From: shahules786 <shahules786@gmail.com>
Date: Thu, 10 Nov 2022 16:01:06 +0530
Subject: [PATCH] add cache

---
 enhancer/models/model.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/enhancer/models/model.py b/enhancer/models/model.py
index c679669..9f285d3 100644
--- a/enhancer/models/model.py
+++ b/enhancer/models/model.py
@@ -18,8 +18,11 @@ from enhancer.inference import Inference
 from enhancer.loss import LOSS_MAP, LossWrapper
 from enhancer.version import __version__
 
-CACHE_DIR = ""
-HF_TORCH_WEIGHTS = ""
+CACHE_DIR = os.getenv(
+    "ENHANCER_CACHE",
+    os.path.expanduser("~/.cache/torch/enhancer"),
+)
+HF_TORCH_WEIGHTS = "pytorch_model.ckpt"
 DEFAULT_DEVICE = "cpu"