import tensorflow as tf

		KK = tf.matmul(K, K, transpose_b=True)
		K_trace = tf.expand_dims(tf.expand_dims(tf.trace(KK), -1), -1)
		K_loss = tf.reduce_mean(tf.abs(KK / K_trace - tf.eye(2))) 

		loss_total_gen = crit_gen + rep_loss + K_loss

	gen_var = model.get_gen_vars()
	dis_var = model.dis.trainable_variables
	grads = tape.gradient([loss_total_gen, crit_dis], [gen_var, dis_var])
	return grads, [crit_dis, crit_gen, rep_loss, K_loss]

reader = datareader.DataReader(16)
model = network.RepNet()
optim = tf.optimizers.Adam(0.0001, 0.5)
saver = M.Saver(model)
saver.restore('./model/')

MAXITER = 10000

bar = tqdm(range(MAXITER+1))
for i in bar:
	batch = reader.get_next()
	grads, lss = grad_loss(batch, model)

	gen_var = model.get_gen_vars()
	dis_var = model.dis.trainable_variables
