Hello,
I added autocast to this a2c model. It is %30 faster, and it looks like it trains well, but can’t be sure if I am doing something wrong here. I would be grateful if you can point it out any mistakes. Thanks.
‘’’
def train_agent():
state = env.reset()

state = torch.from_numpy(state).float().to(device)

for i_episode in range(episode_start, 4000000):       # 1000000

    log_probs = []
    values    = []
    rewards   = []
    masks     = []
    entropy = 0

    states = torch.zeros([num_steps, 720, env.shape[1]]).to(device)       

    actions = []
    dones = []
    final_reward = []

    for step in range(num_steps):

        states[step] = state

        with autocast():

            with torch.no_grad():                  

                value, action, action_log_prob = model.act(state.unsqueeze(1))

        action = action.item()

        next_state, reward, done, info = env.step(action)            

        rewards.append(reward)

        dones.append(0 if done else 1)

        actions.append(action)

        if done:

            final_reward.append(info)            

        state = next_state

        state = torch.from_numpy(state).float().to(device)

    with autocast():

        with torch.no_grad():

            next_value, _ = model(state.unsqueeze(1))            

        

        returns = compute_returns(next_value, rewards, dones)

        returns = torch.cat(returns)

        actions = torch.tensor(actions).unsqueeze(1).to(device)

        states = states.permute(1, 0, 2)

        values, action_log_probs, dist_entropy = model.evaluate_actions(states, actions)            

        advantages = returns - values

        value_loss = advantages.pow(2).mean()

        action_loss = -(advantages.detach() * action_log_probs.unsqueeze(1)).mean()
        
        value_loss_coef = 0.5

        entropy_coef = 0.01

        loss = (value_loss * value_loss_coef + action_loss - dist_entropy * entropy_coef)

    scaler.scale(loss).backward()        

    scaler.unscale_(optimizer)

    max_grad_norm = 0.5

    nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)        

    scaler.step(optimizer)

    scaler.update()

    optimizer.zero_grad()

‘’’