(x, drop_prob: float = 0.0, training: bool = False)
| 33 | |
| 34 | |
| 35 | def drop_path(x, drop_prob: float = 0.0, training: bool = False): |
| 36 | if drop_prob == 0.0 or not training: |
| 37 | return x |
| 38 | keep_prob = 1 - drop_prob |
| 39 | shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets |
| 40 | random_tensor = x.new_empty(shape).bernoulli_(keep_prob) |
| 41 | if keep_prob > 0.0: |
| 42 | random_tensor.div_(keep_prob) |
| 43 | output = x * random_tensor |
| 44 | return output |
| 45 | |
| 46 | |
| 47 | class DropPath(nn.Module): |