Skip to content

Commit f815a8c

Browse files
committed
polish migrate_all_to_cpu
1 parent 16bffae commit f815a8c

File tree

2 files changed

+4
-4
lines changed

2 files changed

+4
-4
lines changed

python/jittor/__init__.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
# file 'LICENSE.txt', which is part of this source code package.
1010
# ***************************************************************
1111

12-
__version__ = '1.3.4.0'
12+
__version__ = '1.3.4.1'
1313
from jittor_utils import lock
1414
with lock.lock_scope():
1515
ori_int = int

python/jittor/src/var_holder.cc

+3-3
Original file line numberDiff line numberDiff line change
@@ -240,9 +240,9 @@ void migrate_all_to_cpu() {
240240
#ifdef HAS_CUDA
241241
for (auto vh : hold_vars) {
242242
auto v = vh->var;
243-
if (v->_outputs.size()) continue;
244-
if (v->allocator && !v->allocator->is_cuda())
245-
migrate_to_gpu(v, cpu_allocator);
243+
// if (v->_outputs.size()) continue;
244+
if (v->allocator && v->mem_ptr && !v->allocator->is_cuda())
245+
migrate_to_cpu(v, cpu_allocator);
246246
}
247247
#endif
248248
}

0 commit comments

Comments
 (0)