[BUG] natoms can not be obtained in calc_model_devi when atype is a list

[wangzyphysics]

Bug summary

deepmd-kit/deepmd/infer/model_devi.py

Lines 309 to 324 in b875ea8

	energies = []
	forces = []
	virials = []
	natom = atype.shape[-1]
	for dp in models:
	ret = dp.eval(
	coord,
	box,
	atype,
	fparam=fparam,
	aparam=aparam,
	mixed_type=mixed_type,
	)
	energies.append(ret[0] / natom)
	forces.append(ret[1])
	virials.append(ret[2] / natom)

in line 312, when atype is a List, an error will occur.

DeePMD-kit Version

#3480

TensorFlow Version

Backend: PyTorch PT ver: v2.0.0.post200-gc263bd43e8e

How did you download the software?

pip

Input Files, Running Commands, Error Log, etc.

POSCAR_1 file:

 CALYPSO
1.0
  -1.72119715062831        1.72119754046866        1.72119741052188     
   1.72119741052189       -1.72119728057511        1.72119741052188
   1.72119741052189        1.72119728057511       -1.72119741052188
 Mg
           1
Direct
  0.000000000000000E+000  0.000000000000000E+000  0.000000000000000E+000

error log:

root@bohrium-11265-1111017:~/run/dpgen-nv83j/iter-000000--collect-run-calypso-000000-0/workdir/caly_task.000000# python test.py 
To get the best performance, it is recommended to adjust the number of threads by setting the environment variables OMP_NUM_THREADS, DP_INTRA_OP_PARALLELISM_THREADS, and DP_INTER_OP_PARALLELISM_THREADS. See https://deepmd.rtfd.io/parallelism/ for more information.
Traceback (most recent call last):
  File "/root/run/dpgen-nv83j/iter-000000--collect-run-calypso-000000-0/workdir/caly_task.000000/test.py", line 12, in <module>
    graphs = [DeepPot("/root/run/model-ckpt-new.pt")] * 4
  File "/opt/re_e_bias_new/deepmd-kit/deepmd/infer/deep_eval.py", line 324, in __init__
    self.deep_eval = DeepEvalBackend(
  File "/opt/re_e_bias_new/deepmd-kit/deepmd/pt/infer/deep_eval.py", line 104, in __init__
    assert not self.multi_task, "multitask mode currently not supported!"
AssertionError: multitask mode currently not supported!
root@bohrium-11265-1111017:~/run/dpgen-nv83j/iter-000000--collect-run-calypso-000000-0/workdir/caly_task.000000# python test.py 
To get the best performance, it is recommended to adjust the number of threads by setting the environment variables OMP_NUM_THREADS, DP_INTRA_OP_PARALLELISM_THREADS, and DP_INTER_OP_PARALLELISM_THREADS. See https://deepmd.rtfd.io/parallelism/ for more information.
You can use the environment variable DP_INFER_BATCH_SIZE tocontrol the inference batch size (nframes * natoms). The default value is 1024.
Traceback (most recent call last):
  File "/root/run/dpgen-nv83j/iter-000000--collect-run-calypso-000000-0/workdir/caly_task.000000/test.py", line 13, in <module>
    devi = calc_model_devi(coord, cell, atype, graphs)
  File "/opt/re_e_bias_new/deepmd-kit/deepmd/infer/model_devi.py", line 307, in calc_model_devi
    natom = atype.shape[-1]
AttributeError: 'list' object has no attribute 'shape'

test.py

root@bohrium-11265-1111017:~/run/dpgen-nv83j/iter-000000--collect-run-calypso-000000-0/workdir/caly_task.000000# cat test.py 
from deepmd.infer import DeepPot
from deepmd.infer import calc_model_devi
from ase.io import read
import numpy as np

atoms = read("POSCAR_1")
pbc = np.all(atoms.get_pbc())
coord = atoms.get_positions().reshape(1, -1)
cell = atoms.get_cell().array.reshape(1, -1) if pbc else None
type_map = ["Mg", "Al"]
atype = [type_map.index(atom.symbol) for atom in atoms]
graphs = [DeepPot("/root/run/dpgen-nv83j/finetune--run-train-0000/workdir/task.0000/model.ckpt.pt")] * 4
devi = calc_model_devi(coord, cell, atype, graphs)

Steps to Reproduce

python test.py

Further Information, Files, and Links

No response