Usage of loss which is not CrossEntropy

[dedeswim]

Hi,

Thank you very much for this implementation!

I noticed that the only allowed loss is CrossEntropyLoss. As I'd need your module to compute Influence Functions on other kinds of losses, how about adding the possibility to use a different loss, perhaps passed as param in the config dict?

I was thinking about the following changes (to avoid disruptive changes to the API):

pytorch_influence_functions/pytorch_influence_functions/influence_function.py

Lines 8 to 9 in 4df5d2e

	def s_test(z_test, t_test, model, z_loader, gpu=-1, damp=0.01, scale=25.0,
	recursion_depth=5000):

to

def s_test(z_test, t_test, model, z_loader, gpu=-1, damp=0.01, scale=25.0,
           recursion_depth=5000, loss_fn=compute_loss):

And then change

pytorch_influence_functions/pytorch_influence_functions/influence_function.py

Line 45 in 4df5d2e

loss = calc_loss(y, t)

to

loss = loss_fn(y, t)

With something analogous for:

pytorch_influence_functions/pytorch_influence_functions/influence_function.py

Line 75 in 4df5d2e

def grad_z(z, t, model, gpu=-1):

It would be obviously necessary to propagate the loss_fn params to the functions that call grad_z and s_test.

I'd be glad to create a new branch and open a pull request with this change!

Thanks :)