Custom loss function for Mutual Information

I want to create a custom loss function which will calculate the mutual information between two training datasets.

For an example,

x= dataset_1
y= dataset_2

MI = mutual_information(x,y)

How can I do that in pytorch? Thank you so much in advanced.

Define mutual information on datasets. The information theory definition of mutual information is on two RVs.

how did you solve this?