replace transformers.Conv1D with MemSaveLinear

plutonium-239 · Apr 22, 2024 · dd48758 · dd48758
1 parent bba6a11
commit dd48758
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/memsave_torch/nn/__init__.py b/memsave_torch/nn/__init__.py
@@ -7,6 +7,7 @@
 """
 
 import torch.nn as nn
+import transformers
 
 from memsave_torch.nn.BatchNorm import MemSaveBatchNorm2d
 from memsave_torch.nn.Conv1d import MemSaveConv1d
@@ -56,7 +57,7 @@ def convert_to_memory_saving(
     layers = [
         {
             "allowed": linear,
-            "cls": nn.Linear,
+            "cls": (nn.Linear, transformers.Conv1D),
             "convert_fn": MemSaveLinear.from_nn_Linear,
         },
         {"allowed": relu, "cls": nn.ReLU, "convert_fn": MemSaveReLU.from_nn_ReLU},