search

jacobkimmel / pytorch_modelsize

Estimates the size of a PyTorch model in memory
MIT License
351 stars 49 forks source link

Half precision implementation #4

Open je-santos opened 4 years ago

je-santos commented 4 years ago

Hi Jacob,

I tried to estimate the size of my model that uses half precision without any luck. I changed line 34 to HalfTensor, but got the following message:

RuntimeError: "unfolded3d_copy_cpu" not implemented for 'Half'

Do you know how can I achieve this?

Thanks for your great work.

-Javier

jacobkimmel commented 4 years ago

I haven't worked with the HalfTensor structure before, so I'm guessing here. The only copy happening in this tool is copying some_tensor.size() to CPU as a numpy array. Are even the size() results stored as HalfTensor?

If so, swapping out .size() for .shape in the get_output_sizes() method might do the trick.