code: def train(model, optimizer, train_loader, device=None):
      device = device or torch.device("cpu")
      model.train()
      for batch_idx, (data, target) in enumerate(train_loader):
          if batch_idx * len(data) > EPOCH_SIZE:
              return
          data, target = data.to(device), target.to(device)
          optimizer.zero_grad()
          output = model(data)
          loss = F.nll_loss(output, target)
          loss.backward()
          optimizer.step()
  
file_path: ray/doc/source/tune/doc_code/pytorch_optuna.py
imports: import os
  from filelock import FileLock
  import torch.nn
  import torch.nn.functional
  from torchvision import datasets, transforms
  import torch
  from ray import train, tune
  from ray.tune.search.optuna import OptunaSearch
label_desc: Use image node_id 1 for calling a global function with example usage: train(model, optimizer, train_loader, device) without return types