torch_brain.nnΒΆ Embeddings Layers Embedding InfiniteVocabEmbedding RotaryTimeEmbedding SinusoidalTimeEmbedding Transformer modules FeedForward RotaryCrossAttention RotarySelfAttention Losses Loss MSELoss CrossEntropyLoss MallowDistanceLoss Readout Layers MultitaskReadout prepare_for_multitask_readout()