import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, gamma=0.99):
        super(Model, self).__init__()
        self.gamma = gamma

    def forward(self, q_values, actions, rewards, next_q_values, dones):
        curr_q = q_values.gather(1, actions.unsqueeze(1)).squeeze(1)
        next_q_max = next_q_values.max(1)[0]
        target = rewards + self.gamma * next_q_max * (1.0 - dones)
        loss = (curr_q - target) ** 2
        return loss.mean()

batch_size = 1024
num_actions = 6

def get_inputs():
    q_values = torch.randn(batch_size, num_actions, requires_grad=True)
    actions = torch.randint(0, num_actions, (batch_size,))
    rewards = torch.randn(batch_size)
    next_q_values = torch.randn(batch_size, num_actions)
    dones = torch.zeros(batch_size)  # float 0.0 or 1.0
    return [q_values, actions, rewards, next_q_values, dones]

def get_init_inputs():
    return [0.99]