FredZhang7
/

google-safesearch-mini-tfjs

Image Classification

Model card Files Files and versions

FredZhang7 commited on Dec 27, 2022

Commit

6fcab6a

·

1 Parent(s): 25e5000

Upload model

Files changed (1) hide show

Model.py +1 -1

Model.py CHANGED Viewed

@@ -113,5 +113,5 @@ class InceptionV3ModelForImageClassification(PreTrainedModel):
 			out, aux = self(image)
 			if print_tensor:
 				print(out)
-			_, predicted = torch.max(out.data, 1)
 		return self.config.classes[str(predicted.item())]

 			out, aux = self(image)
 			if print_tensor:
 				print(out)
+			_, predicted = torch.max(out.logits, 1)
 		return self.config.classes[str(predicted.item())]