e3fcbb639e115e8afe9600bd06aee81acfda6704,reagent/training/world_model/seq2reward_trainer.py,Seq2RewardTrainer,get_loss,#Seq2RewardTrainer#Any#,57

Before Change

        gamma_mask = torch.Tensor(
            [[gamma ** i for i in range(seq_len)] for _ in range(batch_size)]
        ).transpose(0, 1)
        target_acc_reward = torch.sum(target_rewards * gamma_mask, ).unsqueeze(1)
        // make sure the prediction and target tensors have the same size
        // the size should both be (BATCH_SIZE, 1) in this case.
        assert (

After Change

            .to(training_batch.reward.device)
        )

        target_acc_rewards = torch.cumsum(training_batch.reward * gamma_mask, dim=0)
        target_acc_reward = target_acc_rewards[
            valid_reward_len - 1, torch.arange(batch_size)
        ].unsqueeze(1)

        // make sure the prediction and target tensors have the same size
        // the size should both be (BATCH_SIZE, 1) in this case.
        assert (

Instances