def METHOD_NAME(self):
    def foo(x):
        for i in range(3):
            torch.jit._logging.add_stat_value('foo', 1)
        return x + 1.0
    traced = torch.jit.trace(foo, torch.rand(3, 4))
    logger = torch.jit._logging.LockingLogger()
    logger.set_aggregation_type('foo', torch.jit._logging.AggregationType.AVG)
    old_logger = torch.jit._logging.set_logger(logger)
    try:
        traced(torch.rand(3, 4))
        self.assertEqual(logger.get_counter_val('foo'), 1)
    finally:
        torch.jit._logging.set_logger(old_logger)