code: def fit_one_epoch(
      model,
      train_loader,
      batch_transforms,
      optimizer,
      scheduler,
      amp=False,
  ):
      if amp:
          scaler = torch.cuda.amp.GradScaler()
  
      model.train()
      # Iterate over the batches of the dataset
      pbar = tqdm(train_loader, position=1)
      for images, targets in pbar:
          if torch.cuda.is_available():
              images = images.cuda()
          images = batch_transforms(images)
  
          train_loss = model(images, targets)["loss"]
  
          optimizer.zero_grad()
          if amp:
              with torch.cuda.amp.autocast():
                  train_loss = model(images, targets)["loss"]
              scaler.scale(train_loss).backward()
              # Gradient clipping
              scaler.unscale_(optimizer)
              torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
              # Update the params
              scaler.step(optimizer)
              scaler.update()
          else:
              train_loss = model(images, targets)["loss"]
              train_loss.backward()
              torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
              optimizer.step()
  
          scheduler.step()
  
          pbar.set_description(f"Training loss: {train_loss.item():.6}")
  
imports: import os
  import datetime
  import hashlib
  import logging
  import multiprocessing
  import time
  from pathlib import Path
  import numpy
  import torch
  import wandb
  from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR, OneCycleLR
  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
  from torchvision.transforms.v2 import Compose, GaussianBlur, Normalize, RandomGrayscale, RandomPerspective, RandomPhotometricDistort
  from tqdm.auto import tqdm
  from doctr import transforms
  from doctr.datasets import VOCABS, RecognitionDataset, WordGenerator
  from doctr.models import login_to_hub, push_to_hf_hub, recognition
  from doctr.utils.metrics import TextMatch
  from utils import EarlyStopper, plot_recorder, plot_samples
label_desc: Use image node_id 2 for calling a global function with example usage: fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp) without return types