From 757e6af0af2f8de6da696976f3110cd70085adad Mon Sep 17 00:00:00 2001 From: Volpeon Date: Tue, 7 Feb 2023 20:56:37 +0100 Subject: Fix Lora memory usage --- training/strategy/dreambooth.py | 2 -- training/strategy/lora.py | 4 +--- training/strategy/ti.py | 2 -- 3 files changed, 1 insertion(+), 7 deletions(-) (limited to 'training/strategy') diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py index b4c77f3..8aaed3a 100644 --- a/training/strategy/dreambooth.py +++ b/training/strategy/dreambooth.py @@ -1,4 +1,3 @@ -from contextlib import nullcontext from typing import Optional from functools import partial from contextlib import contextmanager, nullcontext @@ -6,7 +5,6 @@ from pathlib import Path import itertools import torch -import torch.nn as nn from torch.utils.data import DataLoader from accelerate import Accelerator diff --git a/training/strategy/lora.py b/training/strategy/lora.py index 88d1824..92abaa6 100644 --- a/training/strategy/lora.py +++ b/training/strategy/lora.py @@ -1,11 +1,9 @@ -from contextlib import nullcontext from typing import Optional from functools import partial -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from pathlib import Path import torch -import torch.nn as nn from torch.utils.data import DataLoader from accelerate import Accelerator diff --git a/training/strategy/ti.py b/training/strategy/ti.py index d306f18..da2b81c 100644 --- a/training/strategy/ti.py +++ b/training/strategy/ti.py @@ -1,11 +1,9 @@ -from contextlib import nullcontext from typing import Optional from functools import partial from contextlib import contextmanager, nullcontext from pathlib import Path import torch -import torch.nn as nn from torch.utils.data import DataLoader from accelerate import Accelerator -- cgit v1.2.3-54-g00ecf