[dattri.trak] Improving memory efficiency when calculating TRAK

[KurisuTheAmadeus]

Current Implementation of TRAK calculates final result using

    return (running_xinv_XTX_XT @ self.Q.diag().to(self.device)).T

and the introduction of diag matrix creates large memory overhead, especially in the case when number of training data points is large.
A possible fix is as follows

        if train_dataloader is not None:
            return (running_xinv_XTX_XT * running_Q.to(self.device).unsqueeze(0)).T
        return (running_xinv_XTX_XT * self.Q.to(self.device).unsqueeze(0)).T

and I have verified the equivalency using two different setups.