Enable model.eval() usage for YOLO class (#17754)

This commit is contained in:
Laughing 2024-11-25 18:00:24 +08:00 committed by GitHub
parent 65c797a898
commit 6a762564c8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1126,3 +1126,20 @@ class Model(nn.Module):
description of the expected behavior and structure. description of the expected behavior and structure.
""" """
raise NotImplementedError("Please provide task map for your model!") raise NotImplementedError("Please provide task map for your model!")
def eval(self):
"""
Sets the model to evaluation mode.
This method changes the model's mode to evaluation, which affects layers like dropout and batch normalization
that behave differently during training and evaluation.
Returns:
(Model): The model instance with evaluation mode set.
Examples:
>> model = YOLO("yolo11n.pt")
>> model.eval()
"""
self.model.eval()
return self