def apply_compilation(model, optimizer):
    model.encoder = torch.compile(model.encoder, fullgraph=True)
    model.decoder = torch.compile(model.decoder, fullgraph=True)
    model.loss_function = torch.compile(model.loss_function, fullgraph=True)
    optimizer.step = torch.compile(optimizer.step)
    return model, optimizer