code: def test_train_eval(self):
      class TestModel(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.dropout = torch.nn.Dropout(p=1.0)
  
          def forward(self, x):
              x = x.float().mean()
              x = self.dropout(x)  # dropout
              if self.training:
                  x += 100  # add
              else:
                  x *= 0  # mul
              x -= 0  # sub
              return x
  
      model = TestModel()
  
      train_return_nodes = ["dropout", "add", "sub"]
      eval_return_nodes = ["dropout", "mul", "sub"]
  
      def checks(model, mode):
          with torch.no_grad():
              out = model(torch.ones(10, 10))
          if mode == "train":
              # Check that dropout is respected
              assert out["dropout"].item() == 0
              # Check that control flow dependent on training_mode is respected
              assert out["sub"].item() == 100
              assert "add" in out
              assert "mul" not in out
          elif mode == "eval":
              # Check that dropout is respected
              assert out["dropout"].item() == 1
              # Check that control flow dependent on training_mode is respected
              assert out["sub"].item() == 0
              assert "mul" in out
              assert "add" not in out
  
      # Starting from train mode
      model.train()
      fx_model = self._create_feature_extractor(
          model,
          train_return_nodes=train_return_nodes,
          eval_return_nodes=eval_return_nodes,
      )
      # Check that the models stay in their original training state
      assert model.training
      assert fx_model.training
      # Check outputs
      checks(fx_model, "train")
      # Check outputs after switching to eval mode
      fx_model.eval()
      checks(fx_model, "eval")
  
      # Starting from eval mode
      model.eval()
      fx_model = self._create_feature_extractor(
          model,
          train_return_nodes=train_return_nodes,
          eval_return_nodes=eval_return_nodes,
      )
      # Check that the models stay in their original training state
      assert not model.training
      assert not fx_model.training
      # Check outputs
      checks(fx_model, "eval")
      # Check outputs after switching to train mode
      fx_model.train()
      checks(fx_model, "train")
  
imports: import random
  from itertools import chain
  from typing import Mapping, Sequence
  import pytest
  import torch
  from common_utils import set_rng_seed
  from torchvision import models
  from torchvision.models._utils import IntermediateLayerGetter
  from torchvision.models.detection.backbone_utils import BackboneWithFPN, mobilenet_backbone, resnet_fpn_backbone
  from torchvision.models.feature_extraction import create_feature_extractor, get_graph_node_names
label_desc: Use image node_id 8 for calling the TestFxFeatureExtraction obj's underlying member method code with example usage: obj.test_train_eval() and returns: x