From 6fbffb7c62d332a834791e976f8cde40b7636c63 Mon Sep 17 00:00:00 2001 From: Dun Liang Date: Fri, 6 Jan 2023 20:48:04 +0800 Subject: [PATCH] jtorch compatible --- python/jittor/__init__.py | 2 +- python/jittor/misc.py | 2 +- python/jittor/nn.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/python/jittor/__init__.py b/python/jittor/__init__.py index 229da570..8d6472bf 100644 --- a/python/jittor/__init__.py +++ b/python/jittor/__init__.py @@ -9,7 +9,7 @@ # file 'LICENSE.txt', which is part of this source code package. # *************************************************************** -__version__ = '1.3.6.9' +__version__ = '1.3.6.10' from jittor_utils import lock with lock.lock_scope(): ori_int = int diff --git a/python/jittor/misc.py b/python/jittor/misc.py index 8e417c3b..fdf3389b 100644 --- a/python/jittor/misc.py +++ b/python/jittor/misc.py @@ -291,7 +291,7 @@ def median(x,dim=None,keepdim=False, keepdims=False): if dim is None: x = x.reshape(-1) dim=0 - _,x = x.argsort(dim) + _,x = jt.argsort(x, dim) slices = [slice(None) for i in range(dim-1)] k = (x.shape[dim]-1)//2 if keepdim: diff --git a/python/jittor/nn.py b/python/jittor/nn.py index ff414535..3146fe21 100644 --- a/python/jittor/nn.py +++ b/python/jittor/nn.py @@ -498,9 +498,9 @@ def softmax(x, dim=None, log=False): return code_softmax.softmax_v1(x, log) if dim is None: dim = () if log: - a = x-x.max(dim, keepdims=True) + a = x - jt.max(x, dim, keepdims=True) return a - a.exp().sum(dim, keepdims=True).log() - x = (x-x.max(dim, keepdims=True)).exp() + x = (x - jt.max(x, dim, keepdims=True)).exp() return x / x.sum(dim, keepdims=True) jt.Var.softmax = softmax