Mercurial > repos > goeckslab > image_learner

diff image_workflow.py @ 21:d5c582cf74bc draft default tip
planemo upload for repository https://github.com/goeckslab/gleam.git commit eed8c1e1d99a8a0c8f3a6bfdd8af48a5bfa19444
author: goeckslab
date: Tue, 20 Jan 2026 01:25:35 +0000
parents: 64872c48a21f
--- a/image_workflow.py	Tue Jan 06 15:35:11 2026 +0000
+++ b/image_workflow.py	Tue Jan 20 01:25:35 2026 +0000
@@ -367,6 +367,17 @@
         self.args.output_dir.mkdir(parents=True, exist_ok=True)
 
         try:
+            if getattr(self.args, "torch_home", None):
+                torch_home = Path(self.args.torch_home).expanduser().resolve()
+                torch_home.mkdir(parents=True, exist_ok=True)
+                os.environ["TORCH_HOME"] = str(torch_home)
+                try:
+                    import torch
+
+                    torch.hub.set_dir(str(torch_home))
+                except Exception as exc:
+                    logger.warning("Unable to set Torch Hub cache dir: %s", exc)
+
             self._create_temp_dirs()
             self._extract_images()
             csv_path, split_cfg, split_info = self._prepare_data()