I have this code that works and does what I want it to do. However, I would also need the Sequential container to contain not only the linear layers, but also the one hot encoding. How could I achieve this? Thanks!
import torch
import torch.nn as nn
import torch.nn.functional as F
preprocessing_layers = {}
## Input
obs_space = torch.Tensor([0, 1,2,3])
obs_space = obs_space.to(torch.long)
## Sequential
layers = [torch.nn.Linear(obs_space.size(-1), 5)]
preprocessing_layers["direction"] = nn.Sequential(*layers)
## One hot
obs_space = F.one_hot(obs_space,4)
obs_space = obs_space.to(torch.float32)
## Forward
a = preprocessing_layers["direction"](obs_space)