From a3d774a0d6c41801b2393c2a83fa0381c69ae523 Mon Sep 17 00:00:00 2001 From: Harry Baker Date: Sun, 14 May 2023 23:38:46 +0100 Subject: [PATCH] Removed `torch.compile` support --- minerva/trainer.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/minerva/trainer.py b/minerva/trainer.py index e3f9c9ccb..dd66f53e1 100644 --- a/minerva/trainer.py +++ b/minerva/trainer.py @@ -39,7 +39,6 @@ # IMPORTS # ===================================================================================================================== import os -import warnings from contextlib import nullcontext from pathlib import Path from typing import ( @@ -62,7 +61,6 @@ from alive_progress import alive_bar, alive_it from inputimeout import TimeoutOccurred, inputimeout from nptyping import Int, NDArray -from packaging.version import Version from torch import Tensor from torch.nn.modules import Module from torch.nn.parallel import DistributedDataParallel as DDP @@ -416,16 +414,6 @@ def __init__( ) self.model = MinervaDataParallel(self.model, DDP, device_ids=[gpu]) - # Wraps the model in `torch.compile` to speed up computation time. - # Python 3.11+ is not yet supported though, hence the exception clause. - if Version(torch.__version__) > Version("2.0.0"): # pragma: no cover - try: - _compiled_model = torch.compile(self.model) # type: ignore[attr-defined] - assert isinstance(_compiled_model, (MinervaModel, MinervaDataParallel)) - self.model = _compiled_model - except RuntimeError as err: - warnings.warn(str(err)) - def init_wandb_metrics(self) -> None: """Setups up separate step counters for :mod:`wandb` logging of train, val, etc.""" if isinstance(self.writer, Run):