(x, drop_prob: float = 0.0, training: bool = False)
| 12 | |
| 13 | |
| 14 | def drop_path(x, drop_prob: float = 0.0, training: bool = False): |
| 15 | if drop_prob == 0.0 or not training: |
| 16 | return x |
| 17 | keep_prob = 1 - drop_prob |
| 18 | shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets |
| 19 | random_tensor = x.new_empty(shape).bernoulli_(keep_prob) |
| 20 | if keep_prob > 0.0: |
| 21 | random_tensor.div_(keep_prob) |
| 22 | output = x * random_tensor |
| 23 | return output |
| 24 | |
| 25 | |
| 26 | class DropPath(nn.Module): |