def setup_keys(self, pretrained_model):
if self.key_type != 'random':
n = 1 if self.key_type == 'image' else 20 # any number will do
key_x, x_inds = passport_generator.get_key(self.passport_data, n)
key_x = key_x.to(self.device)
key_y, y_inds = passport_generator.get_key(self.passport_data, n)
key_y = key_y.to(self.device)
passport_generator.set_key(pretrained_model, self.model,
key_x, key_y)
def reset_passport():
print('Reset passport')
x, y = get_passport(passport_data, device)
passport_generator.set_key(pretrained_model, model, x, y)