Pr 293 (#294)

* Add enable_flash/cueq attributes to SevenNetModel class

* Add enable_flash/cueq in SevenNetModel class.

* chore

---------

Co-authored-by: hswoo369 <hswoo369@gmail.com>

Author: alphalm4

sevenn/torchsim.py

@@ -74,6 +74,8 @@ def __init__(
         *,  # force remaining arguments to be keyword-only
         modal: str | None = None,
         neighbor_list_fn: Callable | None = None,
+        enable_flash: bool = False,
+        enable_cueq: bool = False,
         device: torch.device | str = 'auto',
         dtype: torch.dtype = torch.float32,
     ) -> None:
@@ -90,6 +92,8 @@ def __init__(
             modal (str | None): modal (fidelity) if given model is multi-modal model.
                 for 7net-mf-ompa, it should be one of 'mpa' (MPtrj + sAlex) or
                 'omat24' (OMat24).
+            enable_cueq (bool): Enable cuEquivariance backend.
+            enable_flash (bool): Enable flashTP backend.
             neighbor_list_fn (Callable): Neighbor list function to use.
                 Default is torch_nl_linked_cell.
             device (torch.device | str): Device to run the model on
@@ -126,7 +130,10 @@ def __init__(
 
         if isinstance(model, (str, Path)):
             cp = load_checkpoint(model)
-            model = cp.build_model()
+            model = cp.build_model(
+                enable_flash=enable_flash,
+                enable_cueq=enable_cueq,
+            )
 
         _validate(model, modal)