def METHOD_NAME(args):
    if args.cuda:
        torch.set_default_tensor_type("torch.cuda.FloatTensor")
    else:
        torch.set_default_tensor_type("torch.FloatTensor")
    if args.method == "all":
        for method in ["prob", "logprob", "gradient", "marginal", "map", "sample"]:
            args.method = method
            METHOD_NAME(args)
        return
    print("Plate size  Time per iteration of {} (ms)".format(args.method))
    fn = globals()["jit_{}".format(args.method)]
    equation = args.equation
    plates = args.plates
    inputs, outputs = equation.split("->")
    inputs = inputs.split(",")
    # Vary all plate sizes at the same time.
    for plate_size in range(8, 1 + args.max_plate_size, 8):
        operands = []
        for dims in inputs:
            shape = torch.Size(
                [plate_size if d in plates else args.dim_size for d in dims]
            )
            operands.append((torch.empty(shape).uniform_() + 0.5).requires_grad_())
        time = time_fn(
            fn, equation, *operands, plates=plates, modulo_total=True, iters=args.iters
        )
        print(
            "{: <11s} {:0.4g}".format(
                "{} ** {}".format(plate_size, len(args.plates)), time * 1e3
            )
        )