Changes

    return torch.nn.CrossEntropyLoss(reduction='none')

def _make_tf_shapes_consistent(output, labels):

  """Try to make inputs have the same shape by adding dimensions of size 1."""

  import tensorflow as tf

  flow_layers = [

      tfb.RealNVP(

          num_masked=2,

          num_masked=1,

          shift_and_log_scale_fn=tfb.real_nvp_default_template(

              hidden_layers=[8, 8]))

  ]

  # 3D Multivariate Gaussian base distribution

  nf = NormalizingFlow(

      base_distribution=tfd.MultivariateNormalDiag(loc=[0., 0., 0.]),

      base_distribution=tfd.MultivariateNormalDiag(loc=[0., 0.]),

      flow_layers=flow_layers)

  nfm = NormalizingFlowModel(nf)

  # Must be float32 for RealNVP

  dataset = NumpyDataset(

      X=np.random.rand(5, 3).astype(np.float32),

      y=np.random.rand(5,),

      ids=np.arange(5))

  target_distribution = tfd.MultivariateNormalDiag(loc=[1., 0.])

  dataset = NumpyDataset(X=target_distribution.sample(96))

  # Tests a simple flow of one RealNVP layer.

  X = nfm.flow.sample()

  x1 = tf.zeros([3])

  x1 = tf.zeros([2])

  x2 = dataset.X[0]

  # log likelihoods should be negative

  # # Fit model

  final = nfm.fit(dataset, nb_epoch=5)

  print(final)

  assert final > 0