GB0246

Graph-Break Type

Short name describing what triggered the graph break

Attempted to use torch.nn.functional.one_hot with data-dependent output shape

Context

Values or code snippet captured at the break point

args={args}, kwargs={kwargs}

Explanation

Explanation of why the graph break was triggered

Dynamo does not support this.

Hints

Hints on how to resolve the graph break

• Explicitly set the num_classes param of the function call
• torch.nn.functional.one_hot to something other than -1.

Additional Information

This graph break happens when torch.nn.functional.one_hot is called without a num_classes argument (or with num_classes=-1). In that case the size of the output’s last dimension depends on the maximum value in the input tensor, so the output shape depends on the tensor’s contents, which torch.compile cannot determine while building the compiled function.

Example code that causes the graph break is:

import torch
import torch.nn.functional as F


@torch.compile(fullgraph=True)
def fn(labels):
    return F.one_hot(labels)


fn(torch.tensor([0, 1, 2]))

Pass an explicit num_classes so the output shape is known when the function is compiled. In most models the number of classes is a fixed hyperparameter, so this is also the more robust way to write the code:

import torch
import torch.nn.functional as F


@torch.compile(fullgraph=True)
def fn(labels):
    return F.one_hot(labels, num_classes=10)


fn(torch.tensor([0, 1, 2]))

Click here to add Additional Info

Back to Registry