Merge pull request NVIDIA#904 from ksivaman/sparsity

Fixed variable name
HarryZJ · Jul 16, 2020 · 4027bcb · 4027bcb
2 parents 1ff54b8 + 59995c7
commit 4027bcb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/apex/contrib/sparsity/asp.py b/apex/contrib/sparsity/asp.py
@@ -102,7 +102,7 @@ def add_sparse_attributes(module_name, module):
                         print("[ASP] Sparsifying %s::%s of size=%s and type=%s for sparsity" % (module_name, p_name, str(p.size()), str(p.dtype)))
 
                     mask = torch.ones_like(p).bool()
-                    buffname = name.split(".")[-1] # buffer names cannot contain "."
+                    buffname = p_name.split(".")[-1] # buffer names cannot contain "."
                     module.register_buffer('__%s_mma_mask' % buffname, mask)
                     if allow_recompute_mask:
                         pruned = torch.zeros_like(p).cpu()