Revert "fix: unnecessary memory alloc for large systems"

This reverts commit d8269cb.

Author: iProzd

deepmd/dpmodel/infer/deep_eval.py

@@ -366,7 +366,7 @@ def _eval_model(
                     out = np.full(shape, np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION)
                 results.append(out)
             else:
-                shape = self._get_output_shape(odef, 1, 1)
+                shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
                     np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION)
                 )  # this is kinda hacky

deepmd/jax/infer/deep_eval.py

@@ -389,7 +389,7 @@ def _eval_model(
                     out = np.full(shape, np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION)
                 results.append(out)
             else:
-                shape = self._get_output_shape(odef, 1, 1)
+                shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
                     np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION)
                 )  # this is kinda hacky

deepmd/pd/infer/deep_eval.py

@@ -577,7 +577,7 @@ def _eval_model(
                     out = out.numpy()
                 results.append(out)
             else:
-                shape = self._get_output_shape(odef, 1, 1)
+                shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
                     np.full(np.abs(shape), np.nan, dtype=prec)
                 )  # this is kinda hacky
@@ -657,7 +657,7 @@ def _eval_model_spin(
                 out = batch_output[pd_name].reshape(shape).detach().cpu().numpy()
                 results.append(out)
             else:
-                shape = self._get_output_shape(odef, 1, 1)
+                shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
                     np.full(
                         np.abs(shape),

deepmd/pt/infer/deep_eval.py

@@ -528,7 +528,7 @@ def _eval_model(
                 out = batch_output[pt_name].reshape(shape).detach().cpu().numpy()
                 results.append(out)
             else:
-                shape = self._get_output_shape(odef, 1, 1)
+                shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
                     np.full(np.abs(shape), np.nan, dtype=prec)
                 )  # this is kinda hacky
@@ -608,7 +608,7 @@ def _eval_model_spin(
                 out = batch_output[pt_name].reshape(shape).detach().cpu().numpy()
                 results.append(out)
             else:
-                shape = self._get_output_shape(odef, 1, 1)
+                shape = self._get_output_shape(odef, nframes, natoms)
                 results.append(
                     np.full(
                         np.abs(shape),