StochasticDepth

A stochastic depth layer that wraps a given Component.

StochasticDepth(nn.Linear(128, 256), p=0.2)

Forward

(x: jaxtyping.Float[Tensor, ' *shape']) -> jaxtyping.Float[Tensor, ' *shape']