Multi-gpu compatibility: torch-summary for 'cuda:1' and so on

[adithyavis]

Models in GPU devices other than 'cuda:0' can be summarized after this edit.

An edit was made in the way the kwarg 'device' accepted values. Instead of only accepting string values, the current edit allows for torch.device objects and (legacy) integer values to be accepted. summary(model, (1,28,28),device="cuda:2" summary(model, (1,28,28),device=torch.device("cuda")

[daniel347x]

Alas, it looks like this is a dead GitHub project. This PR has been in for close to a year and half, but no cigar.

Therefore this can't be used during production runs which require control over the GPU on which the code runs.