Documentiation issue line 71,72 optimal_transport.py

[noorschepersoveral-dotcom]

Dear Tong,

First of all, thank you for this page, it has proven quite useful. While figuring out how the code works, I think I have stumbled upon a small documentation mistake. In line 71 and 72 of optimal_transport.py, x1 is described as Tensor, shape (bs, *dim) represents the source minibatch. I think it should be the target minibatch.

Thank you!

Best regards

[atong01]

conditional-flow-matching/torchcfm/optimal_transport.py

Lines 63 to 74 in 66b236b

	def get_map(self, x0, x1):
	"""Compute the OT plan (wrt squared Euclidean cost) between a source and a target
	minibatch.
	Parameters
	----------
	x0 : Tensor, shape (bs, *dim)
	represents the source minibatch
	x1 : Tensor, shape (bs, *dim)
	represents the source minibatch
	Returns

Agreed