code: def _train(self, epochs: int, theta: float, device: str, **kwargs):
      # Define model
      model = self.model.to(device)
      model.epochs = epochs
      batch_size = self.train_data.batch_size
  
      # TODO: SELECT THE N_CLASSES OUTSIDE THE OPERATION
      label_injector = LabelsInjector(datasets.MNIST.classes)
  
      progressive_dataset = ProgressiveTrainingDataset(
          (
              label_injector.inject_train(x, y)
              for x, y in self.train_data
          )
      )
      progressive_dataloader = torch.utils.data.DataLoader(
          progressive_dataset, batch_size=2 * batch_size, shuffle=False
      )
  
      model.train()
      model.progressive_train(progressive_dataloader, theta)
  
      model.eval()
      correct = 0
      with torch.no_grad():
          for data, target in self.test_data:
              input_data = label_injector.inject_eval(data)
              input_data = input_data.to(device)
              target = target.to(device)
              input_shapes = input_data.shape[:-1]
              input_data = input_data.reshape(-1, input_data.shape[-1])
              _, prob = model.positive_eval(input_data, theta)
              prob = prob.reshape(*input_shapes)
              pred = prob.argmax(dim=1)
              correct += (pred == target).float().sum().item()
      if isinstance(correct, torch.Tensor):
          correct = correct.item()
      self.logger.info(
          "Test set: Accuracy: {}/{} ({:.0f}%)".format(
              correct,
              len(self.test_data.dataset),
              100.0 * correct / len(self.test_data.dataset),
          )
      )
  
imports: from abc import ABC, abstractmethod
  import torch
  from nebullvm.operations.base import Operation
  from nebullvm.operations.fetch_operations.local import FetchModelFromLocal
  from torch.utils.data import DataLoader
  from torchvision import datasets
  from forward_forward.operations.data import VOCABULARY
  from forward_forward.operations.fetch_operations import FetchTrainingDataFromLocal
  from forward_forward.utils.labels import LabelsInjector
  from forward_forward.utils.modules import FCNetFFProgressive
  from forward_forward.utils.utils import ProgressiveTrainingDataset, compute_perplexity
label_desc: Use image node_id 1 for calling the ForwardForwardTrainer obj's underlying member method code with example usage: obj._train(epochs, theta, device) without return types