plotting not working, losses fixed?

losses.py

@@ -21,7 +21,7 @@ def d_logistic(real_images, generator, discriminator, z_dim, conditions):
     d_loss += tf.math.softplus(-real_scores)
 
 
-    return fake_images, d_loss
+    return d_loss
 
 
 def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions):
@@ -46,7 +46,7 @@ def d_logistic_r1_reg(real_images, generator, discriminator, z_dim, conditions):
     real_grads = r1_tape.gradient(real_loss, real_images)
     r1_penalty = tf.reduce_sum(tf.math.square(real_grads), axis=[1, 2, 3])
     r1_penalty = tf.expand_dims(r1_penalty, axis=1)
-    return fake_images, d_loss, r1_penalty
+    return d_loss, r1_penalty
 
 
 def g_logistic_non_saturating(real_images, generator, discriminator, z_dim, conditions):
@@ -104,7 +104,7 @@ F = np.load("custom_utils/smpl_faces.npy")
 
 def _nearest_neighbour(V, B):
 	tree = cKDTree(B.numpy())
-	return tree.query(V.numpy(), n_jobs=-1)[1]
+	return tree.query(V.numpy(), workers=-1)[1]
 
 @tf.function
 def _tf_nn_parallel(V, B):
@@ -117,13 +117,14 @@ def collision_loss(V, B, mask, thr=.004):
     # compute correspondences
     idx = _tf_nn_parallel(V, B)
     N, M = V.shape[:2]
-    _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    # _idx = tf.tile(np.array(list(range(N)), np.int32)[:,None], [1, M])
+    _idx = tf.ones((N, M), dtype=tf.int32) * np.array(list(range(N)), np.int32)[:,None]
     idx = tf.stack((_idx, idx), axis=-1)
     # compute loss
     D = V - tf.gather_nd(B, idx)
     b_vn = tf.gather_nd(B_vn, idx)
     dot = tf.einsum('abc,abc->ab', D, b_vn)
-    loss = tf.reduce_sum(tf.minimum(dot * mask - thr, 0) ** 2)
+    loss = tf.reduce_sum((tf.minimum(dot - thr, 0)  * mask) ** 2)
 
     return loss
 
@@ -150,7 +151,6 @@ def edge_loss(x, y, mask):
     return loss
 
 
-
 class PairwiseDistance(tf.keras.layers.Layer):
     def __init__(self, **kwargs):
         super(PairwiseDistance, self).__init__(**kwargs)
@@ -191,7 +191,6 @@ def collision_penalty(va, vb, mask, eps=1e-3):
 if __name__ == '__main__':
     V = tf.ones((4, 14475, 3))
     B = tf.zeros((4, 14475, 3))
-    B_vn = vertex_normals(B, tf.tile(F[None], [B.shape[0], 1, 1]))
-    loss_1 = collision_loss(V, B, thr=.004, stop_gradient=False)
-    loss_2 = collision_penalty(V, B, B_vn)
+    m = tf.ones((4, 14475))
+    loss_1 = collision_loss(V, B, m, thr=.004)
     print()