import set_gpus
set_gpus.set("2")

from config.trecgnet_config import TRECGNET_CONFIG
from config.secret_config import SECRET_CONFIG
from net.trecgnet import TrecgNet
import torch
import os


if __name__ == '__main__':
    net_config = TRECGNET_CONFIG()
    # print(net_config.CHECKPOINT_DIR)
    net_config.TB_LOGOUT = False
    net = TrecgNet(net_config)

    # 加载模型
    net.load_state_dict(torch.load(os.path.join(SECRET_CONFIG.PROJECT_PATH, "checkpoints/default/trecgnet.pth")))

    # 保存Encoder
    net.save_encoder("rgb_encoder/encoder.pth")
    print("Encoder Saved!")
