probly.method.het_nets.torch.drop_in_place_het_layer

probly.method.het_nets.torch.drop_in_place_het_layer(obj: nn.Linear, state: State) → tuple[nn.Module, State]

Replace the last linear layer with a HeteroscedasticLayer.