I am trying to use this architecture:
class Net(BaseFeaturesExtractor): def __init__(self, observation_space: gym.spaces.Box, features_dim: int = 512):