The problem is fixed by not setting end_points as a class variable and returning the entire dict. I suspect the problem is along the lines of https://github.com/pytorch/pytorch/issues/16532 5. Not sure though where the tensor is going out of scope and triggering a recursive deletion of the rest of the graph though.