Save Transformer module in root when training with ST

#50
by tomaarsen HF staff - opened
Files changed (1) hide show
  1. custom_st.py +2 -0
custom_st.py CHANGED
@@ -33,6 +33,8 @@ class Transformer(nn.Module):
33
  None, then model_name_or_path is used
34
  """
35
 
 
 
36
  def __init__(
37
  self,
38
  model_name_or_path: str,
 
33
  None, then model_name_or_path is used
34
  """
35
 
36
+ save_in_root: bool = True
37
+
38
  def __init__(
39
  self,
40
  model_name_or_path: str,