encoder = torch.compile(encoder, mode="reduce-overhead")
decoder = torch.compile(decoder, dynamic=True)