Custom torch model with BatchNorm

How severe does this issue affect your experience of using Ray?

  • Low: It annoys or frustrates me for a moment.

Hi!
I’m currently planning to build a custom ResNet feature extractor to be used for your (awesome) AlphaZero implementation. Looking into the docs and rllib/examples, it seems like BatchNorm layers are only possible using Tensorflow? Or do you only provide TF examples because torch doesn’t need special treatment when using BatchNorm layers?

Kind regards,
Sebastian

Nvm., I saw that you provide a PyTorch BatchNorm model inside examples/batch_norm_model.py! Thanks! :slight_smile: