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