elif hps.output_dist == 'gaussian':
      assert np.issubdtype(
          datasets[hps.dataset_names[0]]['train_data'].dtype, float), \
          "Data dtype must be float for gaussian output dsitribution"
      data_dtype = tf.float32