elif obs_type == 'Conv1d':
return schema_noise_mask_observation(obs_schema['Element'], obs_sample)
elif obs_type == 'Conv2d':
return schema_noise_mask_observation(obs_schema['Element'], obs_sample)
Hey it should work if you add these two lines to that method. We just need to forward the conv2d struct to the element that it wraps.