From 2825ba2f2030b2fd3e841aad416a4fd28d67615a Mon Sep 17 00:00:00 2001 From: Volpeon Date: Sun, 4 Dec 2022 09:24:37 +0100 Subject: Update --- pipelines/util.py | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 pipelines/util.py (limited to 'pipelines/util.py') diff --git a/pipelines/util.py b/pipelines/util.py new file mode 100644 index 0000000..661dbee --- /dev/null +++ b/pipelines/util.py @@ -0,0 +1,9 @@ +import torch + + +def set_use_memory_efficient_attention_xformers(module: torch.nn.Module, valid: bool) -> None: + if hasattr(module, "set_use_memory_efficient_attention_xformers"): + module.set_use_memory_efficient_attention_xformers(valid) + + for child in module.children(): + set_use_memory_efficient_attention_xformers(child, valid) -- cgit v1.2.3-54-g00ecf