my first attemp at visualizing with tensorboard.
built a convnet with nn.module class -

class ThreeLayerConvNet2(nn.Module):
    def __init__(self, in_channel, channel_1, channel_2, num_classes):
        super().__init__()
       
        pad1 = 2
        pad2 = 1
        ker1 = (5,5)
        ker2 = (3,3)
        feature_map1_size = (32 + 2 * pad1 - ker1[0] + 1, 32 + 2 * pad1 - ker1[1] + 1)
        feature_map2_size = (feature_map1_size[0] + 2 * pad2 - ker2[0] + 1, feature_map1_size[1] + 2 * pad2 - ker2[1] + 1)

        flat_size = np.prod(feature_map2_size) * channel_2
        self.conv1 = nn.Conv2d(in_channel, channel_1, kernel_size = ker1, padding = pad1,
        padding_mode = 'zeros')
        nn.init.kaiming_normal_(self.conv1.weight)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(channel_1, channel_2, kernel_size = ker2, padding = pad2,
        padding_mode = 'zeros')
        nn.init.kaiming_normal_(self.conv2.weight)
        self.fc = nn.Linear(flat_size, num_classes)
        #self.batchnorm2d_1 = nn.BatchNorm2d(in_channel)
        #self.batchnorm2d_2 = nn.BatchNorm2d(channel_1)
        #self.batchnorm1d = nn.BatchNorm1d(flat_size)
        pass


    def forward(self, x):
        scores = None
 
        #x = self.batchnorm2d_1(x)
        x = self.conv1(x)
        x = self.relu(x)
        #x = self.batchnorm2d_2(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = flatten(x)
        #x = self.batchnorm1d(x)
        scores = self.fc(x)
        pass


        return scores

num_classes = 10
in_channel = 3
channel_1 = 32
channel_2 = 16
learning_rate = 1e-2

model = ThreeLayerConvNet2(in_channel, channel_1, channel_2, num_classes)

optimizer = optim.SGD(model.parameters(), lr=learning_rate,
                     momentum=0.9, nesterov=True)


then initialized tensorboard -
from torch.utils.tensorboard import SummaryWriter

# default `log_dir` is "runs" - we'll be more specific here
writer = SummaryWriter('cs231n/final_cifar10_convnet')

tried to feed the board a graph like shown in tutorial -
dataiter = iter(loader_train)
images, labels = dataiter.next()
writer.add_graph(ThreeLayerConvNet2,images)

and got the following error.
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-73-49d61804c4c3> in <module>
     1 dataiter = iter(loader_train)
     2 images, labels = dataiter.next()
----> 3 writer.add_graph(ThreeLayerConvNet2)

~\.conda\envs\torch_env\lib\site-packages\torch\utils\tensorboard\writer.py in add_graph(self, model, input_to_model, verbose)
   705         if hasattr(model, 'forward'):
   706             # A valid PyTorch model should have a 'forward' method
--> 707             self._get_file_writer().add_graph(graph(model, input_to_model, verbose))
   708         else:
   709             # Caffe2 models do not have the 'forward' method

~\.conda\envs\torch_env\lib\site-packages\torch\utils\tensorboard\_pytorch_graph.py in graph(model, args, verbose)
   281         processing.
   282     """
--> 283     with torch.onnx.set_training(model, False):  # TODO: move outside of torch.onnx?
   284         try:
   285             trace = torch.jit.trace(model, args)

~\.conda\envs\torch_env\lib\contextlib.py in __enter__(self)
    79     def __enter__(self):
    80         try:
---> 81             return next(self.gen)
    82         except StopIteration:
    83             raise RuntimeError("generator didn't yield") from None

~\.conda\envs\torch_env\lib\site-packages\torch\onnx\utils.py in set_training(model, mode)
    36         yield
    37         return
---> 38     old_mode = model.training
    39     if old_mode != mode:
    40         model.train(mode)

AttributeError: type object 'ThreeLayerConvNet2' has no attribute 'training'