TypeError in larp_ar.py: 'NoneType' object is not subscriptable when mask parameter is None

[abbasiReza]

Training fails with TypeError: 'NoneType' object is not subscriptable at line 393 in models/larp_ar.py when the mask parameter is not provided during forward pass.

Error Traceback

Traceback (most recent call last):
  File "/workspace/NAR/trainers/larp_ar_fp_trainer.py", line 221, in _iter_step
    logits, loss = self.forward_ar_model(z, c)
  File "/workspace/NAR/trainers/larp_ar_fp_trainer.py", line 201, in forward_ar_model
    logits, loss = self.model_ddp(cond_idx=c, idx=input_tokens, targets=z)
  File "/workspace/NAR/models/larp_ar.py", line 393, in forward
    mask = mask[:, None].to(h.device)
           ~~~~^^^^^^^^^
TypeError: 'NoneType' object is not subscriptable

Root Cause In models/larp_ar.py, the forward() method declares mask as an optional parameter (defaults to None):

def forward(
    self, 
    idx: Optional[torch.Tensor],
    cond_idx: Optional[torch.Tensor],
    input_pos: Optional[torch.Tensor] = None, 
    targets: Optional[torch.Tensor] = None,
    mask: Optional[torch.Tensor] = None,  # <-- Optional parameter
    valid: Optional[torch.Tensor] = None,
    coordinate: list = None,
    next_coordinate: list = None,
):

However, at line 393, the code unconditionally tries to index mask without checking if it's None:

# Line 393
mask = mask[:, None].to(h.device)  # Crashes when mask is None

The larp_ar_fp_trainer.py calls the model without providing the mask parameter:

# Line 201 in larp_ar_fp_trainer.py
logits, loss = self.model_ddp(cond_idx=c, idx=input_tokens, targets=z)
# No mask parameter passed

Additional Issue If a None check is added in larp_ar.py, another assertion fails later in the same file:

File "/workspace/NAR/models/larp_ar.py", line 442, in torch_dynamo_resume_in_forward_at_442
    assert seq_len - 1 == T * H * W