docstr: Call this method to monkey patch optimizer step function so that masks can be applied to
  gradients and weights during training.
  You must call init_model_for_pruning(...) before calling init_optimizer_for_pruning(...)
docstr_tok: ["Call","this","method","to","monkey","patch","optimizer","step","function","so","that","masks","can","be","applied","to","gradients","and","weights","during","training",".","You","must","call","init_model_for_pruning","(","...",")","before","calling","init_optimizer_for_pruning","(","...",")"]
code: def init_optimizer_for_pruning(cls, optimizer):
      """Call this method to monkey patch optimizer step function so that masks can be applied to
      gradients and weights during training.
      You must call init_model_for_pruning(...) before calling init_optimizer_for_pruning(...)
      """
      assert (
          cls.__optimizer is None
      ), "ASP has initialized optimizer already."
      assert (
          cls.__calculate_mask is not None
      ), "Called ASP.init_optimizer_for_pruning before ASP.init_model_for_pruning."
  
      # store pointer to original optimizer step method
      cls.__optimizer = optimizer
      cls.__optimizer.__step = optimizer.step
  
      def __step(opt_self, *args, **kwargs):
          # prune gradients before step method
          with torch.no_grad():
              for (
                  module_name,
                  module,
                  p_name,
                  p,
                  mask,
                  pruned,
              ) in cls.__sparse_parameters:
                  if p.grad is not None:  # thx pjudd
                      p.grad.mul_(mask)
          # call original optimizer step method
          rval = opt_self.__step(*args, **kwargs)
          # prune parameters after step method
          with torch.no_grad():
              for (
                  module_name,
                  module,
                  p_name,
                  p,
                  mask,
                  pruned,
              ) in cls.__sparse_parameters:
                  p.mul_(mask)
          return rval
  
      cls.__optimizer.step = types.MethodType(__step, cls.__optimizer)
  
file_path: apex/apex/contrib/sparsity/asp.py
imports: import types
  import torch
  from .sparse_masklib import create_mask
  from .permutation_lib import Permutation
  import json
  import os
  import string
  import time
label_desc: Use image node_id 3 for calling the ASP obj's underlying member method code with example usage: obj.init_optimizer_for_pruning(cls, optimizer) and returns: rval