search

liuruiyang98 / Jittor-summary

Keras style model.summary() in Jittor, modified from torchsummary.
MIT License
19 stars 0 forks source link

强制summary使用cpu计算 #1

Open cjld opened 3 years ago

cjld commented 3 years ago

https://github.com/liuruiyang98/Jittor-summary/blob/3842c79c4db59b03a5acab1aa8ba23d6085ed53b/jittorsummary/jittorsummary.py#L13

可以在这一行添加一个修饰符@jt.flag_scope(use_cuda=0)来强制使用cpu

liuruiyang98 commented 3 years ago

将jt.int64更换为numpy.int64，现在jittorsummary已支持CUDA版本，通过 device=‘cpu’ 或者 device=‘cuda' 进行控制。