def insert_tensors(self, tensor_dict: TensorDict, data_dict: DataDict,
model_dict: ModelDict,
batch_idx: np.ndarray) -> TensorDict:
source_tensor = tensor_dict.get(self.source_tensor_key)
transformed = self.transform_lambda(source_tensor).to(device)
tensor_dict.set(self.target_tensor_key, transformed)
return tensor_dict
def insert_tensors(self, tensor_dict: TensorDict, data_dict: DataDict,
model_dict: ModelDict,
batch_idx: np.ndarray) -> TensorDict:
tensors = torch.stack([
tensor_dict.get(tensor_key)
for tensor_key in self.source_tensor_keys
])
summed = torch.sum(tensors, dim=0)
tensor_dict.set(self.target_tensor_key, summed)
return tensor_dict
def insert_tensors(self, tensor_dict: TensorDict, data_dict: DataDict,
model_dict: ModelDict,
batch_idx: np.ndarray) -> TensorDict:
model = model_dict.get(self.model_key)
source_tensors = [
tensor_dict.get(key) for key in self.source_tensor_keys
]
result_tensor_list = self.transform_lambda(model, source_tensors)
for target_tensor_key, result_tensor in zip(self.target_tensor_keys,
result_tensor_list):
tensor_dict.set(target_tensor_key, result_tensor)
return tensor_dict
def get_loss(self, tensor_dicts: TensorDict):
return self.weight * torch.mean(
tensor_dicts.get(self.input_tensor_key)**2)