From 7f6b2b451184a59d543b86caad9ce585a909af83 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Niklas=20B=C3=B6hm?= Date: Wed, 6 Dec 2023 11:13:33 +0100 Subject: [PATCH] Fix logic bug in checking ffcv Fixes #5. Now we check for ffcv after we have properly determined if we need it (which is in the tsimcne.fit() function call). --- tsimcne/tsimcne.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tsimcne/tsimcne.py b/tsimcne/tsimcne.py index 81c8021..a8ebf03 100644 --- a/tsimcne/tsimcne.py +++ b/tsimcne/tsimcne.py @@ -341,7 +341,9 @@ def _handle_parameters(self): f"but got {self.freeze_schedule}." ) - if self.use_ffcv: + @staticmethod + def check_ffcv(use_ffcv): + if use_ffcv: try: import ffcv @@ -406,6 +408,7 @@ def fit(self, X: torch.utils.data.Dataset | str): self.use_ffcv = True else: self.use_ffcv = False + self.check_ffcv(self.use_ffcv) train_dl = self.make_dataloader(X, True, None)