lgatr.layers.slim_layers.SlimDropout

class lgatr.layers.slim_layers.SlimDropout(dropout_prob)

Bases: Module

Dropout for vector and scalar features.

For vector features the same dropout mask is applied to all four components of each vector.

Parameters:

dropout_prob (float) – Dropout probability.

forward(vectors, scalars)

Apply dropout.

Parameters:

• vectors (Tensor) – Lorentz vectors of shape (..., 4, v_channels).

• scalars (Tensor) – Scalar features of shape (..., s_channels).

Return type:

tuple[Tensor, Tensor]

Returns:

• outputs_v – Lorentz vectors with dropout, same shape as vectors.

• outputs_s – Scalar features with dropout, same shape as scalars.