text: import torch
  import torch.nn as nn
  from torch.autograd import Variable
  
  
  class ShakeDropFunction(torch.autograd.Function):
  
      @staticmethod
      def forward(ctx, x, training=True, p_drop=0.5, alpha_range=None):
          if alpha_range is None:
              alpha_range = [-1, 1]
          if training:
              gate = torch.cuda.FloatTensor([0]).bernoulli_(1 - p_drop)
              ctx.save_for_backward(gate)
              if gate.item() == 0:
                  alpha = torch.cuda.FloatTensor(x.size(0)).uniform_(*alpha_range)
                  alpha = alpha.view(alpha.size(0), 1, 1, 1).expand_as(x)
                  return alpha * x
              else:
                  return x
          else:
              return (1 - p_drop) * x
  
      @staticmethod
      def backward(ctx, grad_output):
          gate = ctx.saved_tensors[0]
          if gate.item() == 0:
              beta = torch.cuda.FloatTensor(grad_output.size(0)).uniform_(0, 1)
              beta = beta.view(beta.size(0), 1, 1, 1).expand_as(grad_output)
              beta = Variable(beta)
              return beta * grad_output, None, None, None
          else:
              return grad_output, None, None, None
  
  
  class ShakeDrop(nn.Module):
  
      def __init__(self, p_drop=0.5, alpha_range=None):
          super(ShakeDrop, self).__init__()
          if alpha_range is None:
              alpha_range = [-1, 1]
          self.p_drop = p_drop
          self.alpha_range = alpha_range
  
      def forward(self, x):
          return ShakeDropFunction.apply(x, self.training, self.p_drop, self.alpha_range)
  