[Dynamo] 'NoneType' object is not subscriptable from torchrec (bad error message)
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

UPDATE: Actually this is user error, the patch causes an attribute to no longer exist, but Dynamo gives a bad error message instead of correctly noticing user code problem

Steps to reproduce:

pytorch version: https://github.com/pytorch/pytorch/commit/796ce672296c9ae8d7387b18403810aa2f1048a1
torchrec version: 005727ef06cdc808aa7d263ab4f2837938a77ed2

1. Get working torchrec install (you need fbgemm-gpu and torchrec, install by source; explicitly uninstall fbgemm-gpu-nightly and torchrec-nightly if they exist)
2. Patch torchrec with https://gist.github.com/ezyang/d1a15bb53e0cfc6e05c4e6b57a49bc85
3. For ease of diagnosis, patch pytorch with https://github.com/pytorch/pytorch/pull/107683 (but I don't think it is strictly necessary)
4. Run `TORCH_LOGS=+dynamo MASTER_ADDR=127.0.0.1 MASTER_PORT=29501 RANK=0 LOCAL_RANK=0 WORLD_SIZE=1 python train_dlrm.py`

Stack trace

```
Traceback (most recent call last):
  File "/data/users/ezyang/a/torchrec/examples/golden_training/train_dlrm.py", line 164, in <module>
    main()
  File "/data/users/ezyang/a/pytorch/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/examples/golden_training/train_dlrm.py", line 52, in main
    train()
  File "/data/users/ezyang/a/torchrec/examples/golden_training/train_dlrm.py", line 157, in train
    print(model(next(train_iterator).to(device)))
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/distributed/model_parallel.py", line 266, in forward
    return self._dmp_wrapped_module(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/parallel/distributed.py", line 1519, in forward
    else self._run_ddp_forward(*inputs, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/parallel/distributed.py", line 1355, in _run_ddp_forward
    return self.module(*inputs, **kwargs)  # type: ignore[index]
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/eval_frame.py", line 333, in _fn
    return fn(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/eval_frame.py", line 490, in catch_errors
    return hijacked_callback(frame, cache_size, hooks, frame_state)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 626, in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 133, in _fn
    return fn(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 379, in _convert_frame_assert
    return _compile(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 571, in _compile
    raise InternalTorchDynamoError(str(e)).with_traceback(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 554, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/utils.py", line 181, in time_wrapper
    r = func(*args, **kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 476, in compile_inner
    out_code = transform_code_object(code, transform)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/bytecode_transformation.py", line 1028, in transform_code_object
    transformations(instructions, code_options)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/convert_frame.py", line 443, in transform
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2074, in run
    super().run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/nn_module.py", line 331, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1155, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars.items)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 307, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/nn_module.py", line 331, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1155, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars.items)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 307, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/nn_module.py", line 716, in call_function
    ).call_function(tx, [self] + list(args), kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1155, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars.items)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 307, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1155, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars.items)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 307, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/nn_module.py", line 716, in call_function
    ).call_function(tx, [self] + list(args), kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1155, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars.items)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 307, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1115, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/nn_module.py", line 716, in call_function
    ).call_function(tx, [self] + list(args), kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 1155, in CALL_FUNCTION_EX
    self.call_function(fn, argsvars.items, kwargsvars.items)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 562, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 307, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 261, in call_function
    return super().call_function(tx, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/functions.py", line 90, in call_function
    return tx.inline_user_function_return(
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 598, in inline_user_function_return
    result = InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2179, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 2286, in inline_call_
    tracer.run()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 724, in run
    and self.step()
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 688, in step
    getattr(self, inst.opname)(inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 392, in wrapper
    return inner_fn(self, inst)
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/symbolic_convert.py", line 168, in impl
    self.push(fn_var.call_function(self, self.popn(nargs), {}))
  File "/data/users/ezyang/a/pytorch/torch/_dynamo/variables/builtin.py", line 621, in call_function
    self.as_python_constant()(
torch._dynamo.exc.InternalTorchDynamoError: 'NoneType' object is not subscriptable

from user code:
   File "/data/users/ezyang/a/torchrec/torchrec/models/dlrm.py", line 893, in forward
    logits = self.model(batch.dense_features, batch.sparse_features)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/models/dlrm.py", line 571, in forward
    embedded_sparse = self.sparse_arch(sparse_features)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/models/dlrm.py", line 99, in forward
    sparse_features: KeyedTensor = self.embedding_bag_collection(features)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/distributed/types.py", line 699, in forward
    dist_input = self.input_dist(ctx, *input, **kwargs).wait().wait()
  File "/data/users/ezyang/a/torchrec/torchrec/distributed/embeddingbag.py", line 669, in input_dist
    awaitables.append(input_dist(features_by_shard))
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/distributed/sharding/tw_sharding.py", line 240, in forward
    return self._dist(sparse_features)
  File "/data/users/ezyang/a/pytorch/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/users/ezyang/a/torchrec/torchrec/distributed/dist_data.py", line 454, in forward
    self._splits_cumsum[rank] : self._splits_cumsum[rank + 1]
```

Full debug logs: https://gist.github.com/ezyang/3340a5fa3330e135e13250b3b4246804

### Versions

main

cc @msaroufim @wconstab @bdhirsh @anijain2305