Jittor-summary
Jittor-summary copied to clipboard
强制summary使用cpu计算
https://github.com/liuruiyang98/Jittor-summary/blob/3842c79c4db59b03a5acab1aa8ba23d6085ed53b/jittorsummary/jittorsummary.py#L13
可以在这一行添加一个修饰符@jt.flag_scope(use_cuda=0)
来强制使用cpu
将jt.int64更换为numpy.int64,现在jittorsummary已支持CUDA版本,通过 device=‘cpu’ 或者 device=‘cuda' 进行控制。