| 103 | |
| 104 | |
| 105 | def init_model(model_name): |
| 106 | if model_name == "convnext_small": |
| 107 | return torchvision.models.convnext_small() |
| 108 | if model_name == "densenet161": |
| 109 | return torchvision.models.densenet161() |
| 110 | if model_name == "shufflenet_v2_x1_0": |
| 111 | return torchvision.models.shufflenet_v2_x1_0() |
| 112 | if model_name == "YOLO_NAS_S": |
| 113 | try: |
| 114 | from super_gradients.common.object_names import Models |
| 115 | from super_gradients.training import models |
| 116 | except ImportError: |
| 117 | raise ImportError( |
| 118 | "Please install super-gradients to use the YOLO_NAS_S model." |
| 119 | ) |
| 120 | |
| 121 | return models.get(Models.YOLO_NAS_S, pretrained_weights="coco") |
| 122 | |
| 123 | return None |
| 124 | |
| 125 | |
| 126 | def get_sample_inputs(model_name): |