14 lines
306 B
Python
14 lines
306 B
Python
import torch
|
|
|
|
DTYPE_MAPPING = {
|
|
"float32": torch.float32,
|
|
"float16": torch.float16,
|
|
"bfloat16": torch.bfloat16,
|
|
"float64": torch.float64,
|
|
"uint8": torch.uint8,
|
|
"int8": torch.int8,
|
|
"int16": torch.int16,
|
|
"int32": torch.int32,
|
|
"int64": torch.int64,
|
|
"bool": torch.bool
|
|
} |