Searched refs:device_type (Results 1 – 1 of 1) sorted by relevance
13 static PetscErrorCode EnumToDeviceType(TorchDeviceType device_enum, torch::DeviceType *device_type)… in EnumToDeviceType() argument17 *device_type = torch::kCPU; in EnumToDeviceType()20 *device_type = torch::kXPU; in EnumToDeviceType()23 *device_type = torch::kCUDA; in EnumToDeviceType()26 *device_type = torch::kHIP; in EnumToDeviceType()34 …ic PetscErrorCode PetscMemTypeToDeviceType(PetscMemType mem_type, torch::DeviceType *device_type) { in PetscMemTypeToDeviceType() argument38 *device_type = torch::kCPU; in PetscMemTypeToDeviceType()41 *device_type = torch::kXPU; in PetscMemTypeToDeviceType()44 *device_type = torch::kCUDA; in PetscMemTypeToDeviceType()47 *device_type = torch::kHIP; in PetscMemTypeToDeviceType()