text: import torch
  import torch.nn as nn
  from torch.autograd import Function
  
  
  class HardSigmoid(nn.Module):
  
      def __init__(self):
          super(HardSigmoid, self).__init__()
          self.act = nn.Hardtanh()
  
      def forward(self, x):
          return (self.act(x) + 1.0) / 2.0
  
  
  class RoundFunctionST(Function):
      """Rounds a tensor whose values are in [0, 1] to a tensor with values in {0, 1}"""
  
      @staticmethod
      def forward(ctx, input):
          """Forward pass
  
          Parameters
          ==========
          :param input: input tensor
  
          Returns
          =======
          :return: a tensor which is round(input)"""
  
          # We can cache arbitrary Tensors for use in the backward pass using the
          # save_for_backward method.
          # ctx.save_for_backward(input)
  
          return torch.round(input)
  
      @staticmethod
      def backward(ctx, grad_output):
          """In the backward pass we receive a tensor containing the gradient of the
          loss with respect to the output, and we need to compute the gradient of the
          loss with respect to the input.
  
          Parameters
          ==========
          :param grad_output: tensor that stores the gradients of the loss wrt. output
  
          Returns
          =======
          :return: tensor that stores the gradients of the loss wrt. input"""
  
          # This is a pattern that is very convenient - at the top of backward
          # unpack saved_tensors and initialize all gradients w.r.t. inputs to
          # None. Thanks to the fact that additional trailing Nones are
          # ignored, the return statement is simple even when the function has
          # optional inputs.
          # input, weight, bias = ctx.saved_variables
  
          return grad_output
  
  
  class BernoulliFunctionST(Function):
  
      @staticmethod
      def forward(ctx, input):
  
          return torch.bernoulli(input)
  
      @staticmethod
      def backward(ctx, grad_output):
  
          return grad_output
  
  
  RoundST = RoundFunctionST.apply
  BernoulliST = BernoulliFunctionST.apply
  