[fake tensor] AllenaiLongformerBase - Could not find common device for aten.div.Tensor

[anijain2305]

benchmarks/huggingface.py -dcuda --backend=eager --float32 --train --accuracy --only=AllenaiLongformerBase

Repro

import torch

import torchdynamo
torchdynamo.config.raise_on_assertion_error = True

window_overlap = 256


def fn(x):
    return torch.div(x.size(1), (window_overlap * 2))

x = torch.randn(torch.Size([12, 1024, 64]))
fn(x)

opt_fn = torchdynamo.optimize("eager")(fn)
opt_fn(x)

Error

  File "/scratch/anijain/work/pytorch/torch/_subclasses/fake_tensor.py", line 845, in wrap_meta_outputs_with_default_device_logic
    return tree_map(partial(wrap), r)
  File "/scratch/anijain/work/pytorch/torch/utils/_pytree.py", line 192, in tree_map
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/scratch/anijain/work/pytorch/torch/utils/_pytree.py", line 192, in <listcomp>
    return tree_unflatten([fn(i) for i in flat_args], spec)
  File "/scratch/anijain/work/pytorch/torch/_subclasses/fake_tensor.py", line 857, in wrap
    common_device = FakeTensor._find_common_device(func, args, kwargs)
  File "/scratch/anijain/work/pytorch/torch/_subclasses/fake_tensor.py", line 592, in _find_common_device
    assert common_device is not None, f"Could not find common device for {func}"
AssertionError: Could not find common device for aten.div.Tensor