(self, param_names)
| 130 | return mgx.argument_from_pointer(mgx_shape, tensor.data_ptr()) |
| 131 | |
| 132 | def prealloc_buffers(self, param_names): |
| 133 | for param_name in param_names: |
| 134 | param_shape = self.model.get_parameter_shapes()[param_name] |
| 135 | |
| 136 | type_str, lens = param_shape.type_string(), param_shape.lens() |
| 137 | strides = param_shape.strides() |
| 138 | torch_dtype = self.mgx_to_torch_dtype_dict[type_str] |
| 139 | tensor = torch.empty_strided(lens, |
| 140 | strides, |
| 141 | dtype=torch_dtype, |
| 142 | device=torch.cuda.current_device()) |
| 143 | self.torch_buffers[param_name] = tensor |
| 144 | self.mgx_args[param_name] = self.tensor_to_arg(tensor) |
| 145 | |
| 146 | def get_run_times(self): |
| 147 | return [ |
no test coverage detected