Skip to content

Commit 2ade2cf

Browse files
JackUrbmseeger
authored andcommitted
Deferring import of torch in config to allow faster import (Lightning-AI#2079)
1 parent 1c98f3e commit 2ade2cf

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

litgpt/config.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
from pathlib import Path
66
from typing import Any, List, Literal, Optional, Type, Union
77

8-
import torch
98
import yaml
109
from typing_extensions import Self
1110

@@ -185,6 +184,8 @@ def norm_class(self) -> Type:
185184

186185
from functools import partial
187186

187+
import torch # Torch import is lazy to make config loading faster
188+
188189
if self.norm_class_name == "RMSNorm":
189190
from litgpt.model import RMSNorm
190191

0 commit comments

Comments
 (0)