(self, x)
| 109 | self.drop_prob = drop_prob |
| 110 | |
| 111 | def forward(self, x): |
| 112 | if self.drop_prob == 0.0 or not self.training: |
| 113 | return x |
| 114 | keep_prob = 1 - self.drop_prob |
| 115 | shape = (x.shape[0],) + (1,) * (x.ndim - 1) |
| 116 | random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) |
| 117 | random_tensor.floor_() |
| 118 | output = x.div(keep_prob) * random_tensor |
| 119 | return output |
| 120 | |
| 121 | class GlobalAvgPool2d(nn.Module): |
| 122 | def __init__(self): |