code: def forward(self, x):
      device = x.device
  
      if self.prob == 0.0 or (not self.training):
          return x
  
      keep_mask = (
          torch.FloatTensor(
              (x.shape[0], 1, 1, 1), device=device
          ).uniform_()
          > self.prob
      )
      return x * keep_mask / (1 - self.prob)
  
imports: from functools import partial
  import torch
  from torch import nn, einsum
  from einops import rearrange, repeat
  from einops.layers.torch import Rearrange, Reduce
label_desc: Use image node_id 2 for calling the Dropsample obj's underlying member method code with example usage: obj.forward(x) and returns: unknown, x