From 123977ace3478cbe9bc7adc3ad5bd579afafab22 Mon Sep 17 00:00:00 2001 From: Parmeet Singh Bhatia Date: Wed, 25 May 2022 18:18:25 -0400 Subject: [PATCH] Take TORCH_HOME env variable into account while setting the cache dir --- torchtext/__init__.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/torchtext/__init__.py b/torchtext/__init__.py index 81ba54acd2..331a3fd74a 100644 --- a/torchtext/__init__.py +++ b/torchtext/__init__.py @@ -4,7 +4,11 @@ from torchtext import _extension # noqa: F401 _TEXT_BUCKET = "https://download.pytorch.org/models/text/" -_CACHE_DIR = os.path.expanduser("~/.cache/torch/text") + +_TORCH_HOME = os.getenv("TORCH_HOME") +if _TORCH_HOME is None: + _TORCH_HOME = "~/.cache/torch" #default +_CACHE_DIR = os.path.expanduser(os.path.join(_TORCH_HOME, "text")) from . import data, datasets, experimental, functional, models, nn, transforms, utils, vocab