text: #!/usr/bin/env python3
  # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
  
  from unittest import TestCase
  
  from pytext.common.constants import Stage
  from pytext.metric_reporters.channel import TensorBoardChannel
  from torch import nn, optim
  
  
  class FCModelWithNanAndInfWts(nn.Module):
      """ Simple FC model
      """
  
      def __init__(self):
          super(FCModelWithNanAndInfWts, self).__init__()
          self.fc1 = nn.Linear(10, 5)
          self.fc2 = nn.Linear(5, 3)
          self.fc1.weight.data.fill_(float("NaN"))
          self.fc2.weight.data.fill_(float("Inf"))
  
      def forward(self, x):
          x = self.fc1(x)
          return self.fc2(x)
  
  
  class TensorboardTest(TestCase):
      def test_report_metrics_with_nan(self):
          """ Check that tensorboard channel catches errors when model has
              Inf or NaN weights
          """
          tensorboard_channel = TensorBoardChannel()
          # create simple model and optimizers
          model = FCModelWithNanAndInfWts()
  
          optimizer = optim.SGD(model.parameters(), lr=0.1)
          tensorboard_channel.report(
              stage=Stage.TRAIN,
              epoch=1,
              metrics=0.0,
              model_select_metric=0.0,
              loss=1.0,
              preds=[1],
              targets=[1],
              scores=[1],
              context={},
              meta={},
              model=model,
              optimizer=optimizer,
          )
  