network: IdentityAug

network_kwargs:
  input_shape: null