(
self,
devices: List[Union[str, torch.device]],
device_batch_size: int = 64,
cache_dir: Optional[str] = None,
)
| 50 | |
| 51 | class PointNetClassifier(FeatureExtractor): |
| 52 | def __init__( |
| 53 | self, |
| 54 | devices: List[Union[str, torch.device]], |
| 55 | device_batch_size: int = 64, |
| 56 | cache_dir: Optional[str] = None, |
| 57 | ): |
| 58 | state_dict = load_checkpoint("pointnet", device=torch.device("cpu"), cache_dir=cache_dir)[ |
| 59 | "model_state_dict" |
| 60 | ] |
| 61 | |
| 62 | self.device_batch_size = device_batch_size |
| 63 | self.devices = devices |
| 64 | self.models = [] |
| 65 | for device in devices: |
| 66 | model = get_model(num_class=40, normal_channel=False, width_mult=2) |
| 67 | model.load_state_dict(state_dict) |
| 68 | model.to(device) |
| 69 | model.eval() |
| 70 | self.models.append(model) |
| 71 | |
| 72 | @property |
| 73 | def supports_predictions(self) -> bool: |
nothing calls this directly
no test coverage detected