Rename WrappedGPT to WandaGPT

neuralmagic · Dec 1, 2023 · 563efd7 · 563efd7
1 parent b8d4dff
commit 563efd7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/sparseml/modifiers/pruning/wanda/utils/layer_compressor.py b/src/sparseml/modifiers/pruning/wanda/utils/layer_compressor.py
@@ -19,7 +19,7 @@
 import torch
 from torch.nn import Module
 
-from sparseml.modifiers.pruning.wanda.utils.wrapped_gpt import WrappedGPT
+from sparseml.modifiers.pruning.wanda.utils.wrapped_gpt import WandaGPT
 from sparseml.pytorch.utils.helpers import get_dependency_order
 from sparseml.utils.pytorch.module import get_prunable_layers
 
@@ -76,7 +76,7 @@ def pre_compress_parallel(self, **kwargs) -> Dict:
 
         gpts = {}
         for name in subset:
-            gpts[name] = WrappedGPT(subset[name])
+            gpts[name] = WandaGPT(subset[name])
 
         def add_batch(name):
             def tmp(_, inp, out):
@@ -167,7 +167,7 @@ def sequentially_compress(self, **kwargs):
 
         nsamples = len(self.inputs)
         for name in order:  # create WrappedGPT object for each compressible module
-            gpts = WrappedGPT(subset[name])
+            gpts = WandaGPT(subset[name])
 
             def add_batch(name):
                 def tmp(_, inp, out):

diff --git a/src/sparseml/modifiers/pruning/wanda/utils/wrapped_gpt.py b/src/sparseml/modifiers/pruning/wanda/utils/wrapped_gpt.py
@@ -25,7 +25,7 @@
     transformers = None
     transformers_err = err
 
-__all__ = ["WrappedGPT"]
+__all__ = ["WandaGPT"]
 
 
 DEBUG = False
@@ -35,7 +35,7 @@
 torch.backends.cudnn.allow_tf32 = False
 
 
-class WrappedGPT:
+class WandaGPT:
     """
     Runs Wanda on a single module that contains no sub-modules