Function cond_predicate

intermediate_source/torch_export_tutorial.py:933–941

The conditional statement (aka predicate) passed to ``cond()`` must be one of the following: - ``torch.Tensor`` with a single element - boolean expression NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.

(x)

# We can look at the example called ``cond_predicate``, which has a ``torch.cond`` tag. The example code looks like:

def cond_predicate(x):
    """
    The conditional statement (aka predicate) passed to ``cond()`` must be one of the following:
    - ``torch.Tensor`` with a single element
    - boolean expression
    NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
    """
    pred = x.dim() > 2 and x.shape[2] > 10
    return cond(pred, lambda x: x.cos(), lambda y: y.sin(), [x])

######################################################################
# More generally, ExportDB can be used as a reference when one of the following occurs: