update temp

Author: floatingCatty

dptb/data/interfaces/ham_to_feature.py

@@ -405,6 +405,7 @@ def feature_to_block(data, idp):
                     blocks[block_index] = block.T
                 else:
                     blocks[block_index] += block.T
+                    
     return blocks
 
 

dptb/entrypoints/run.py

@@ -9,7 +9,7 @@
 from dptb.utils.argcheck import normalize_run
 from dptb.utils.tools import j_loader
 from dptb.utils.tools import j_must_have
-from dptb.postprocess.write_ham import write_ham
+from dptb.postprocess.write_block import write_block
 import torch
 import h5py
 
@@ -88,7 +88,7 @@ def run(
 
     elif task=='write_block':
         task = torch.load(init_model, map_location="cpu")["task"]
-        block = write_ham(data=struct_file, AtomicData_options=jdata['AtomicData_options'], model=model, device=jdata["device"])
+        block = write_block(data=struct_file, AtomicData_options=jdata['AtomicData_options'], model=model, device=jdata["device"])
         # write to h5 file, block is a dict, write to a h5 file
         with h5py.File(os.path.join(results_path, task+".h5"), 'w') as fid:
             default_group = fid.create_group("0")

dptb/postprocess/__init__.py

@@ -1,11 +1,11 @@
 from .bandstructure import Band
 from .totbplas import TBPLaS
-from .write_ham import write_ham
+from .write_block import write_block
 
 
 __all__ = [
     Band,
     TBPLaS,
-    write_ham,
+    write_block,
 
 ]

dptb/postprocess/write_ham.py renamed to dptb/postprocess/write_block.py

@@ -13,7 +13,7 @@
 
 log = logging.getLogger(__name__)
 
-def write_ham(
+def write_block(
         data: Union[AtomicData, ase.Atoms, str], 
         model: torch.nn.Module,
         AtomicData_options: dict={},
@@ -35,11 +35,12 @@ def write_ham(
         data = data
 
     data = AtomicData.to_AtomicDataDict(data.to(device))
-    data = model.idp(data)
+    with torch.no_grad():
+        data = model.idp(data)
 
-    # set the kpoint of the AtomicData
-    data = model(data)
-    block = feature_to_block(data=data, idp=model.idp)
+        # set the kpoint of the AtomicData
+        data = model(data)
+        block = feature_to_block(data=data, idp=model.idp)
 
     return block