Привет, вот мой код обучения модели:
```py

# train.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from model import PriceDirectionTransformer
from utils import BTCUSDTDataset
import numpy as np

def train_model(model, dataloader, criterion, optimizer, device):
    model.train()
    epoch_loss = 0
    correct_predictions = 0
    total_predictions = 0
    for sequences, labels in dataloader:
        sequences = sequences.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        outputs = model(sequences)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        _, predicted = torch.max(outputs, 1)
        correct_predictions += (predicted == labels).sum().item()
        total_predictions += labels.size(0)
    accuracy = correct_predictions / total_predictions
    return epoch_loss / len(dataloader), accuracy

def validate_model(model, dataloader, criterion, device):
    model.eval()
    epoch_loss = 0
    correct_predictions = 0
    total_predictions = 0
    with torch.no_grad():
        for sequences, labels in dataloader:
            sequences = sequences.to(device)
            labels = labels.to(device)
            outputs = model(sequences)
            loss = criterion(outputs, labels)
            epoch_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            correct_predictions += (predicted == labels).sum().item()
            total_predictions += labels.size(0)
    accuracy = correct_predictions / total_predictions
    return epoch_loss / len(dataloader), accuracy

if __name__ == "__main__":
    # Hyperparameters
    FEATURE_SIZE = 64
    NUM_LAYERS = 2
    NUM_HEADS = 4
    DROPOUT = 0.1
    BATCH_SIZE = 512
    NUM_EPOCHS = 50
    LEARNING_RATE = 0.001
    WEIGHT_DECAY = 1e-5
    PATIENCE = 5  # Early stopping patience

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

    print(f"Using device: {device}")

    # Load Data
    X_train = np.load('data/X_train.npy')
    y_train = np.load('data/y_train.npy')
    X_val = np.load('data/X_val.npy')
    y_val = np.load('data/y_val.npy')

    # Create Datasets and Dataloaders
    train_dataset = BTCUSDTDataset(X_train, y_train)
    val_dataset = BTCUSDTDataset(X_val, y_val)
    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)

    # Initialize Model
    model = PriceDirectionTransformer(
        feature_size=FEATURE_SIZE,
        num_layers=NUM_LAYERS,
        num_heads=NUM_HEADS,
        dropout=DROPOUT,
    ).to(device)

    # Loss and Optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)

    # Training Loop
    best_val_accuracy = 0
    epochs_without_improvement = 0
    train_losses = []
    val_losses = []
    train_accuracies = []
    val_accuracies = []

    for epoch in range(NUM_EPOCHS):
        train_loss, train_accuracy = train_model(
            model, train_loader, criterion, optimizer, device
        )
        val_loss, val_accuracy = validate_model(
            model, val_loader, criterion, device
        )
        scheduler.step(val_loss)
        train_losses.append(train_loss)
        val_losses.append(val_loss)
        train_accuracies.append(train_accuracy)
        val_accuracies.append(val_accuracy)

        if val_accuracy > best_val_accuracy:
            best_val_accuracy = val_accuracy
            torch.save(model.state_dict(), 'models/best_model.pth')
            epochs_without_improvement = 0
        else:
            epochs_without_improvement += 1
            if epochs_without_improvement >= PATIENCE:
                print(f'Early stopping at epoch {epoch+1}')
                break

        print(
            f'Epoch {epoch+1}/{NUM_EPOCHS}, '
            f'Train Loss: {train_loss:.4f}, Train Acc: {train_accuracy:.4f}, '
            f'Val Loss: {val_loss:.4f}, Val Acc: {val_accuracy:.4f}'
        )
    # Save Final Model
    torch.save(model.state_dict(), 'models/final_model.pth')
    print("Training completed.")

    # Plot training and validation loss and accuracy
    import matplotlib.pyplot as plt
    epochs = range(1, len(train_losses) + 1)
    plt.figure(figsize=(12,5))
    plt.subplot(1,2,1)
    plt.plot(epochs, train_losses, label='Train Loss')
    plt.plot(epochs, val_losses, label='Val Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.title('Loss')

    plt.subplot(1,2,2)
    plt.plot(epochs, train_accuracies, label='Train Acc')
    plt.plot(epochs, val_accuracies, label='Val Acc')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.title('Accuracy')
    plt.show()
```
Можно ли добавить сюда прогрессбар, чтобы можно было видеть прогресс обучения каждой эпохи