code: def train_model(self):
      torch.manual_seed(self.seed)
      device = torch.device("cuda" if self.use_cuda else "cpu")
  
      self.model = Net()
      self.model.to(device)
      optimizer = optim.Adadelta(self.model.parameters(), lr=self.lr)
  
      scheduler = StepLR(optimizer, step_size=1, gamma=self.gamma)
      for epoch in range(1, self.epochs + 1):
          train(
              self.model,
              device,
              self.train_loader,
              optimizer,
              epoch,
              self.log_interval,
              self.dry_run,
          )
          test(self.model, device, self.test_loader)
          scheduler.step()
  
      if self.save_model:
          torch.save(self.model.state_dict(), "mnist_cnn.pt")
  
      self.next(self.end)
  
imports: import os
  from pathlib import Path
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  import wandb
  from metaflow import FlowSpec, Parameter, step
  from torch.optim.lr_scheduler import StepLR
  from torchvision import datasets, transforms
  from wandb.integration.metaflow import wandb_log
label_desc: Use image node_id 4 for calling the WandbPyTorchFlow obj's underlying member method code with example usage: obj.train_model() without return types