forked from ccf-ai-infra/GPUCodeForces
28 lines
734 B
Python
28 lines
734 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super(Model, self).__init__()
|
|
|
|
def forward(self, pred_rewards: torch.Tensor, expert_log_probs: torch.Tensor,
|
|
policy_log_probs: torch.Tensor) -> torch.Tensor:
|
|
expert_loss = -(pred_rewards * expert_log_probs).mean()
|
|
policy_loss = (pred_rewards * policy_log_probs).mean()
|
|
loss = expert_loss + policy_loss
|
|
return loss
|
|
|
|
|
|
batch_size = 32
|
|
|
|
|
|
def get_inputs():
|
|
pred_rewards = torch.randn(batch_size)
|
|
expert_log_probs = torch.randn(batch_size)
|
|
policy_log_probs = torch.randn(batch_size)
|
|
return [pred_rewards, expert_log_probs, policy_log_probs]
|
|
|
|
|
|
def get_init_inputs():
|
|
return [] |