I'm trying to make distilation of Diffusion model called AuraFlow. Here is my code:
```py
import torch
from torch import nn
from torch.optim import Adam
from diffusers import FlowMatchEulerDiscreteScheduler, AuraFlowTransformer2DModel
from torchvision import transforms
from torch.utils.data import DataLoader

teacher_model: AuraFlowTransformer2DModel = AuraFlowTransformer2DModel.from_pretrained(
    'fal/AuraFlow-v0.3',
    torch_dtype=torch.float16,
    variant='fp16',
    load_in_8bit=False,
    subfolder='transformer'
)
teacher_model.eval()

student_model = AuraFlowTransformer2DModel(
    num_mmdit_layers=2, 
    num_single_dit_layers=16, 
    attention_head_dim=128, 
    num_attention_heads=6
)
student_model.train()

from datasets import load_dataset, Dataset

transform = transforms.Compose([
    # transforms.ToPILImage(),
    transforms.Resize((256, 256)),
    # transforms.CenterCrop(256),
    transforms.ToTensor(),
    # transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def preprocess_image(example):
    example['image'] = [transform(img) for img in example['image']]
    return example

train_dataset = load_dataset('gorovuha/ru-image-captions', split='train')
train_dataset.set_transform(preprocess_image)
train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)

# Define the optimizer
optimizer = Adam(student_model.parameters(), lr=1e-4)

# Define the noise scheduler
noise_scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000)

# Define the loss function (e.g., Mean Squared Error)
criterion = nn.MSELoss()

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Move models to the appropriate device (GPU if available)
teacher_model.to(device)
student_model.to(device)

# Training loop
num_epochs = 5  # Adjust the number of epochs as needed

for epoch in range(num_epochs):
    for i, batch in enumerate(train_dataloader):
        # Get the images
        images: torch.Tensor = batch['image']
        images = images.to(device)
        
        # Sample random timesteps for each image in the batch
        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (images.shape[0],), device=device).long()
        
        # Add noise to the images according to the noise scheduler
        noise: torch.Tensor = torch.randn_like(images)
        noisy_images: torch.Tensor = noise_scheduler.scale_noise(images, timesteps, noise)
        
        with torch.autocast(device_type='cuda'):
            with torch.no_grad():
                # Teacher model prediction
                teacher_pred = teacher_model.forward(hidden_states=noisy_images, timestep=timesteps).sample
        
            # Student model prediction
            student_pred: torch.Tensor = student_model.forward(hidden_states=noisy_images, timesteps=timesteps)
            
            # Compute distillation loss
            loss: torch.Tensor = criterion.forward(student_pred, teacher_pred)
        
        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    print(f"Epoch {epoch+1}/{num_epochs} - Loss: {loss.item():.4f}")
```
Can you checks this code for errors? Also my latent I suppose is need to be encoded with VAE (noisy_images or images?)