def __init__(self,
dir_process: str = './',
use_cuda: bool = False,
*args,
**kwargs):
Device(use_cuda)
super().__init__(dir_process, _Plan)
self(HEADER)
def _new_const(self, toward):
if toward.num_type == Exp.BOOL:
return bool(toward.value)
return framework.new_const(toward, Device.get())
def method_randn(toward, args, plan):
args = _float_to_int(args.get_value())
value = _torch.randn(*args, device=Device.get())
return value
def method_tensor(toward, args, plan):
args = _float_to_int(args.get_value())
value = _torch.zeros(*args, device=Device.get())
return value
def _load_binary_raw(cls, path):
raw_np = super()._load_binary_raw(path)
raw_torch = _torch.from_numpy(raw_np).to(Device.get())
return raw_torch