It seems changing the WithItemBuilder class from (PyTorch 1.6.0):
class WithItemBuilder(Builder):
    @staticmethod
    def build_withitem(ctx, item):
        lineno = item.context_expr.lineno
        start = item.context_expr.col_offset
        op_vars = item.optional_vars
        if op_vars:
            end = op_vars.col_offset + len(op_vars.id)
        else:
            end = start + len(item.context_expr.id)

        r = ctx.make_range(lineno, start, end)

        return WithItem(r, build_expr(ctx, item.context_expr), build_expr(ctx, op_vars) if op_vars else None)

to (PyTorch master branch)
class WithItemBuilder(Builder):
    @staticmethod
    def build_withitem(ctx, item):
        lineno = item.context_expr.lineno
        start = item.context_expr.col_offset
        end = start + len(pretty_node_names[ast.With])
        op_vars = item.optional_vars
        r = ctx.make_range(lineno, start, end)

        return WithItem(r, build_expr(ctx, item.context_expr), build_expr(ctx, op_vars) if op_vars else None)

would resolve the issue, however I am not sure if this is really the fix to the problem.