@ptrblck thanks  I was doing it the way  u suggested but the thing is variable returned by function which is sckcit metric function I think is not having require grad as true, so loss.backward  threw exception saying no grad  or grad fn existing, so that is why I had to add this ,it resolved the error but is now always resulting into GPU meomory error with just 2 to 3 iterations of epoch .