ds=false, dont use ds for now

Dockerfile

-FROM nvidia/cuda:11.3.1-runtime-ubuntu20.04
+FROM nvidia/cuda:11.3.1-base-ubuntu20.04
 ENV DEBIAN_FRONTEND=noninteractive
 
 #Install core packages
@@ -26,7 +26,7 @@ RUN pip3 install -e stable-diffusion-private-hypernets/.
 RUN pip3 install https://github.com/crowsonkb/k-diffusion/archive/481677d114f6ea445aa009cf5bd7a9cdee909e47.zip
 RUN pip3 install simplejpeg
 RUN pip3 install min-dalle
-RUN pip3 install https://github.com/microsoft/DeepSpeed/archive/55b7b9e008943b8b93d4903d90b255313bb9d82c.zip
+#RUN pip3 install https://github.com/microsoft/DeepSpeed/archive/55b7b9e008943b8b93d4903d90b255313bb9d82c.zip
 RUN pip3 install https://www.dropbox.com/s/4zq3vq04e79flyh/basedformer-main.zip?dl=1
 
 #Open ports

hydra_node/models.py

@@ -581,5 +581,5 @@ class BasedformerModel(nn.Module):
         prompt = request.prompt
         prompt = self.tokenizer.encode("Input: " + prompt, return_tensors='pt').cuda().long()
         prompt = torch.cat([prompt, torch.tensor([[49527]], dtype=torch.long).cuda()], dim=1)
-        is_safe, corrected = generate(self.model.module, prompt, self.tokenizer, tokens_to_generate=150, ds=True)
+        is_safe, corrected = generate(self.model, prompt, self.tokenizer, tokens_to_generate=150, ds=False)
         return is_safe, corrected
\ No newline at end of file