96f486bd47be8baad365a32fd34d22881fdafa95,reagent/training/ranking/seq2slate_sim_trainer.py,Seq2SlateSimulationTrainer,_simulated_training_input,#Seq2SlateSimulationTrainer#Any#Any#Any#Any#,146

Before Change

        ).detach()
        if slate_reward.ndim == 1:
            logger.warning(f"Slate reward should be 2-D tensor, unsqueezing")
            slate_reward = slate_reward.unsqueeze(1)
        elif slate_reward.ndim != 2:
            raise RuntimeError("Expect slate reward to be 2-D tensor")
        // guard-rail reward prediction range
        reward_clamp = self.parameters.simulation_reward_clamp

After Change

            )

        sim_slate_reward = torch.zeros_like(training_input.slate_reward)
        for name, reward_net in self.reward_name_and_net.items():
            weight = self.sim_param.reward_name_weight[name]
            sr = reward_net(
                training_input.state.float_features,
                training_input.src_seq.float_features,
                sim_tgt_out_seq.float_features,
                training_input.src_src_mask,
                sim_tgt_out_idx,
            ).detach()
            assert sr.ndim == 2, f"Slate reward {name} output should be 2-D tensor"
            sim_slate_reward += weight * sr

        // guard-rail reward prediction range
        reward_clamp = self.sim_param.reward_clamp
        if reward_clamp is not None:
            sim_slate_reward = torch.clamp(
                sim_slate_reward, min=reward_clamp.clamp_min, max=reward_clamp.clamp_max

Instances