code: def train_epoch(model, dataloader, optimizer, e):
      model.train()
      n_examples = 0.0
      accumulated = torch.zeros(1, dtype=torch.float64).cuda()
  
      for batch in dataloader:
          batch = move_to(batch, torch.cuda.current_device())
          speakers = torch.tensor(batch["rest"][2]).cuda()
  
          speaker_logprobs = model(batch)
          loss = F.nll_loss(speaker_logprobs, speakers)
  
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          accumulated += loss.detach().sum()
          n_examples += speakers.size(0)
  
      train_loss = (accumulated / n_examples).item()
      print(f"Epoch {e} | sliding mean train loss {train_loss}")
  
file_path: textlesslib/examples/speaker_probing/train.py
imports: import torch
  import argparse
  from textless.data.quantized_datasets import QuantizedLibriSpeech
  from torch.utils.data import DataLoader
  import torch.nn.functional
  from probes import ContinuousClassifier, DiscreteClassifier, ConstantBaseline
  from textless.data.speech_encoder import SpeechEncoder
  from textless import dispatch_dense_model, dispatch_quantizer
label_desc: Use image node_id 5 for calling a global function with example usage: train_epoch(model, dataloader, optimizer, e) without return types