def METHOD_NAME(self):
    graph_ir = """
    graph(%x: Tensor):
      %half_to_float: bool = prim::Constant[value=1]()
      %dim: int = prim::Constant[value=1]()
      %y = aten::_log_softmax(%x, %dim, %half_to_float)
      return (%y)
    """
    x = torch.randn(5, 2).half().to("cuda")
    self.run_test(graph_ir, (x,))