text: Title: Error in trying to run Example: MCMC Methods for Tall Data
  Body: I try to run the code in 
  
  http://num.pyro.ai/en/stable/examples/covtype.html?highlight=GPU#sphx-glr-download-examples-covtype-py
  
  TypeError: Argument 'None' of type '<class 'NoneType'>' is not a valid JAX type