docstr: Make sure activation size is correct for more complex module
docstr_tok: ["Make","sure","activation","size","is","correct","for","more","complex","module"]
code: def test_activation_size(self) -> None:
      """Make sure activation size is correct for more complex module"""
  
      class TestModule(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.hidden = torch.nn.Linear(10, 5)
              self.relu = torch.nn.ReLU()
              self.out = torch.nn.Linear(5, 3)
              self.softmax = torch.nn.Softmax(dim=3)
  
          def forward(self, x):
              x = self.hidden(x)
              x = self.relu(x)
              x = self.out(x)
              x = self.softmax(x)
              return x
  
      model = TestModule()
      ms = get_module_summary(model, (torch.randn(1, 3, 10, 10),))
      self.assertEqual(ms.module_name, "")
      self.assertEqual(ms.in_size, [1, 3, 10, 10])
      self.assertEqual(ms.out_size, [1, 3, 10, 3])
      self.assertEqual(ms.module_type, "TestModule")
  
      ms_hidden = ms.submodule_summaries["hidden"]
      self.assertEqual(ms_hidden.module_name, "hidden")
      self.assertEqual(ms_hidden.in_size, [1, 3, 10, 10])
      self.assertEqual(ms_hidden.out_size, [1, 3, 10, 5])
      self.assertEqual(ms_hidden.module_type, "Linear")
  
      ms_relu = ms.submodule_summaries["relu"]
      self.assertEqual(ms_relu.module_name, "relu")
      self.assertEqual(ms_relu.in_size, [1, 3, 10, 5])
      self.assertEqual(ms_relu.out_size, [1, 3, 10, 5])
      self.assertEqual(ms_relu.module_type, "ReLU")
  
      ms_out = ms.submodule_summaries["out"]
      self.assertEqual(ms_out.module_name, "out")
      self.assertEqual(ms_out.in_size, [1, 3, 10, 5])
      self.assertEqual(ms_out.out_size, [1, 3, 10, 3])
      self.assertEqual(ms_out.module_type, "Linear")
  
      ms_softmax = ms.submodule_summaries["softmax"]
      self.assertEqual(ms_softmax.module_name, "softmax")
      self.assertEqual(ms_softmax.in_size, [1, 3, 10, 3])
      self.assertEqual(ms_softmax.out_size, [1, 3, 10, 3])
      self.assertEqual(ms_softmax.module_type, "Softmax")
  
imports: import unittest
  from typing import Any, Dict, Optional, Tuple
  import torch
  from torchtnt.utils.module_summary import _get_human_readable_count, get_module_summary, ModuleSummary, prune_module_summary
  from torchvision import models
  from torchvision.models.alexnet import AlexNet_Weights
  from torchvision.models.resnet import ResNet18_Weights
label_desc: Use image node_id 3 for calling the ModuleSummaryTest obj's underlying member method code with example usage: obj.test_activation_size() and returns: x