for input in task.batches(split=split):
input = input.to(device)
output = eval_mygpt(model, input, fixed_len=fixed_len)
- loss = F.cross_entropy(output.transpose(1, 2), input)
+ if args.noncausal_prompt:
+ t = input.size(1) // 2
+ loss = F.cross_entropy(output[:, t:].transpose(1, 2), input[:, t:])
+ else:
+ loss = F.cross_entropy(output.transpose(1, 2), input)
acc_loss += loss.item() * input.size(0)
nb_samples += input.size(0)
if args.noncausal_prompt:
amm_generator = lambda d: torch.logical_and(
torch.arange(d)[None, None, :, None] < torch.arange(d)[None, None, None, :],
- torch.arange(d)[None, None, :, None] >= d // 2,
+ torch.logical_or(
+ torch.arange(d)[None, None, :, None] >= d // 2,
+ torch.arange(d)[None, None, None, :] >= d // 2,
+ ),
)
model = mygpt.MyGPT(
output = eval_mygpt(
model, input, mode=args.oneshot_input, fixed_len=task.height * task.width
)
- loss = F.cross_entropy(output.transpose(1, 2), input)
+ if args.noncausal_prompt:
+ t = input.size(1) // 2
+ loss = F.cross_entropy(output[:, t:].transpose(1, 2), input[:, t:])
+ else:
+ loss = F.cross_entropy(output.transpose(1, 2), input)
acc_train_loss += loss.item() * input.size(0)
nb_train_samples += input.size(0)