code: def __init__(self):
      super().__init__()
      self.M_1_stem = M_1_stem().to("cuda:0")
      self.M_2_stem = M_2_stem().to("cuda:1")
      self.M_1_flatten = torch.nn.Flatten().to("cuda:0")
      self.M_2_flatten = torch.nn.Flatten().to("cuda:1")
      self.M_1_fc1 = torch.nn.Linear(
          out_features=256, in_features=1024
      ).to("cuda:0")
      self.M_2_fc1 = torch.nn.Linear(
          out_features=256, in_features=1024
      ).to("cuda:1")
      self.M_1_fc2 = torch.nn.Linear(
          out_features=10, in_features=256
      ).to("cuda:0")
      self.M_2_fc2 = torch.nn.Linear(
          out_features=10, in_features=256
      ).to("cuda:1")
      self.M_1_softmax = torch.nn.Softmax().to("cuda:0")
      self.M_2_softmax = torch.nn.Softmax().to("cuda:1")
  
imports: from pathlib import Path
  import pytest
  import torch
  import torch.nn
  import torchmetrics
  from torchvision.datasets import MNIST
  from torchvision import transforms
  from pytorch_lightning import seed_everything
  import nni
  from nni.experiment.config import RemoteConfig, RemoteMachineConfig
  from nni.nas.evaluator.pytorch import Lightning, DataLoader
  from nni.nas.execution import SequentialExecutionEngine
  from nni.nas.execution.cgo import CrossGraphOptimization
  from nni.nas.execution.cgo.evaluator import MultiModelLightningModule, MultiModelTrainer
  from nni.nas.execution.cgo.logical_optimizer.logical_plan import LogicalPlan
  from nni.nas.execution.cgo.logical_optimizer.opt_dedup_input import DedupInputOptimizer
  from nni.nas.space import Node, ModelStatus
  from nni.nas.space.pytorch import PytorchGraphModelSpace
  from nni.runtime.trial_command_channel import get_default_trial_command_channel, set_default_trial_command_channel
  from ut.sdk.helper.trial_command_channel import TestHelperTrialCommandChannel
label_desc: Use image node_id 1 to create a new _model_gpu object from inherited base classes: nn with example: obj = _model_gpu()