class_docstr: `StochasticBlock` extends `HybridBlock` to support accumulating loss
  in the forward phase, which is extremely useful in building Bayesian Neural Network,
  where the loss function is composed of a classification loss and a KL loss.
class_docstr_tok: ["`","StochasticBlock","`","extends","`","HybridBlock","`","to","support","accumulating","loss","in","the","forward","phase",",","which","is","extremely","useful","in","building","Bayesian","Neural","Network",",","where","the","loss","function","is","composed","of","a","classification","loss","and","a","KL","loss","."]
docstr: To accumulate loss during the forward phase, one could first decorate
  forward with `StochasticBlock.collectLoss,
  and then collect the loss tensor `x` by calling self.add_loss(x).
  For example, in the following forward function,
  we generate samples from a Gaussian parameterized by `loc` and `scale` and
  accumulate the KL-divergence between it and its prior into the block's loss storage.:
  @StochasticBlock.collectLoss
  def forward(self, loc, scale):
      qz = mgp.Normal(loc, scale)
      # prior
      pz = mgp.Normal(np.zeros_like(loc), np.ones_like(scale))
      self.add_loss(mgp.kl_divergence(qz, pz))
      return qz.sample()
docstr_tok: ["To","accumulate","loss","during","the","forward","phase",",","one","could","first","decorate","forward","with","`","StochasticBlock.collectLoss",",","and","then","collect","the","loss","tensor","`","x","`","by","calling","self.add_loss","(","x",")",".","For","example",",","in","the","following","forward","function",",","we","generate","samples","from","a","Gaussian","parameterized","by","`","loc","`","and","`","scale","`","and","accumulate","the","KL-divergence","between","it","and","its","prior","into","the","block","'s","loss","storage",".",":","@","StochasticBlock.collectLoss","def","forward","(","self",",","loc",",","scale",")",":","qz","=","mgp.Normal","(","loc",",","scale",")","#","prior","pz","=","mgp.Normal","(","np.zeros_like","(","loc",")",",","np.ones_like","(","scale",")",")","self.add_loss","(","mgp.kl_divergence","(","qz",",","pz",")",")","return","qz.sample","(",")"]
code: def collectLoss(func):
      """To accumulate loss during the forward phase, one could first decorate
      forward with `StochasticBlock.collectLoss,
      and then collect the loss tensor `x` by calling self.add_loss(x).
      For example, in the following forward function,
      we generate samples from a Gaussian parameterized by `loc` and `scale` and
      accumulate the KL-divergence between it and its prior into the block's loss storage.:
      @StochasticBlock.collectLoss
      def forward(self, loc, scale):
          qz = mgp.Normal(loc, scale)
          # prior
          pz = mgp.Normal(np.zeros_like(loc), np.ones_like(scale))
          self.add_loss(mgp.kl_divergence(qz, pz))
          return qz.sample()
      """
  
      @wraps(func)
      def inner(self, *args, **kwargs):
          # Loss from forward
          func_out = func(self, *args, **kwargs)
          collected_loss = self._losscache
          self._losscache = []
          self._flag = True
          return (func_out, collected_loss)
  
      return inner
  
file_path: mxnet/python/mxnet/gluon/probability/block/stochastic_block.py
imports: from functools import wraps
  from ...block import HybridBlock
  from ...utils import _indent
label_desc: Use image node_id 3 for calling the StochasticBlock obj's underlying member method code with example usage: obj.collectLoss(func) and returns: inner, func_out,  collected_loss